fix a bug where the spacing would not be transposed in export

MIC-DKFZ · Jun 5, 2024 · 4eb1f5d · 4eb1f5d
1 parent 8fd3580
commit 4eb1f5d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/nnunetv2/evaluation/evaluate_predictions.py b/nnunetv2/evaluation/evaluate_predictions.py
@@ -92,7 +92,6 @@ def compute_metrics(reference_file: str, prediction_file: str, image_reader_writ
     # load images
     seg_ref, seg_ref_dict = image_reader_writer.read_seg(reference_file)
     seg_pred, seg_pred_dict = image_reader_writer.read_seg(prediction_file)
-    # spacing = seg_ref_dict['spacing']
 
     ignore_mask = seg_ref == ignore_label if ignore_label is not None else None
 

diff --git a/nnunetv2/inference/export_prediction.py b/nnunetv2/inference/export_prediction.py
@@ -23,14 +23,15 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(predicted_logits
     torch.set_num_threads(num_threads_torch)
 
     # resample to original shape
+    spacing_transposed = [properties_dict['spacing'][i] for i in plans_manager.transpose_forward]
     current_spacing = configuration_manager.spacing if \
         len(configuration_manager.spacing) == \
         len(properties_dict['shape_after_cropping_and_before_resampling']) else \
-        [properties_dict['spacing'][0], *configuration_manager.spacing]
+        [spacing_transposed[0], *configuration_manager.spacing]
     predicted_logits = configuration_manager.resampling_fn_probabilities(predicted_logits,
                                             properties_dict['shape_after_cropping_and_before_resampling'],
                                             current_spacing,
-                                            properties_dict['spacing'])
+                                            [properties_dict['spacing'][i] for i in plans_manager.transpose_forward])
     # return value of resampling_fn_probabilities can be ndarray or Tensor but that does not matter because
     # apply_inference_nonlin will convert to torch
     predicted_probabilities = label_manager.apply_inference_nonlin(predicted_logits)
@@ -123,13 +124,14 @@ def resample_and_save(predicted: Union[torch.Tensor, np.ndarray], target_shape:
     if isinstance(dataset_json_dict_or_file, str):
         dataset_json_dict_or_file = load_json(dataset_json_dict_or_file)
 
+    spacing_transposed = [properties_dict['spacing'][i] for i in plans_manager.transpose_forward]
     # resample to original shape
     current_spacing = configuration_manager.spacing if \
         len(configuration_manager.spacing) == len(properties_dict['shape_after_cropping_and_before_resampling']) else \
-        [properties_dict['spacing'][0], *configuration_manager.spacing]
+        [spacing_transposed[0], *configuration_manager.spacing]
     target_spacing = configuration_manager.spacing if len(configuration_manager.spacing) == \
         len(properties_dict['shape_after_cropping_and_before_resampling']) else \
-        [properties_dict['spacing'][0], *configuration_manager.spacing]
+        [spacing_transposed[0], *configuration_manager.spacing]
     predicted_array_or_file = configuration_manager.resampling_fn_probabilities(predicted,
                                                                                 target_shape,
                                                                                 current_spacing,