fixes integration test (missing original affine) (Project-MONAI#223)

carlpe · Mar 28, 2020 · 32f8adb · 32f8adb
1 parent 8072028
commit 32f8adb
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 8 deletions.
diff --git a/examples/segmentation_3d/unet_evaluation_dict.py b/examples/segmentation_3d/unet_evaluation_dict.py
@@ -88,7 +88,6 @@
         metric_sum += value.sum().item()
         val_outputs = (val_outputs.sigmoid() >= 0.5).float()
         saver.save_batch(val_outputs, {'filename_or_obj': val_data['img.filename_or_obj'],
-                                       'original_affine': val_data['img.original_affine'],
                                        'affine': val_data['img.affine']})
     metric = metric_sum / metric_count
     print('evaluation metric:', metric)

diff --git a/examples/segmentation_3d_ignite/unet_evaluation_dict.py b/examples/segmentation_3d_ignite/unet_evaluation_dict.py
@@ -100,9 +100,7 @@ def _sliding_window_processor(engine, batch):
 # convert the necessary metadata from batch data
 SegmentationSaver(output_path='tempdir', output_ext='.nii.gz', output_postfix='seg', name='evaluator',
                   batch_transform=lambda batch: {'filename_or_obj': batch['img.filename_or_obj'],
-                                                 'original_affine': batch['img.original_affine'],
-                                                 'affine': batch['img.affine'],
-                                                 },
+                                                 'affine': batch['img.affine']},
                   output_transform=lambda output: predict_segmentation(output[0])).attach(evaluator)
 # the model was trained by "unet_training_dict" exmple
 CheckpointLoader(load_path='./runs/net_checkpoint_50.pth', load_dict={'net': net}).attach(evaluator)

diff --git a/tests/test_integration_segmentation_3d.py b/tests/test_integration_segmentation_3d.py
@@ -203,10 +203,8 @@ def run_inference_test(root_dir, device=torch.device("cuda:0")):
             metric_sum += value.sum().item()
             val_outputs = (val_outputs.sigmoid() >= 0.5).float()
             saver.save_batch(
-                val_outputs, {
-                    'filename_or_obj': val_data['img.filename_or_obj'], 'original_affine':
-                        val_data['img.original_affine'], 'affine': val_data['img.affine']
-                })
+                val_outputs,
+                {'filename_or_obj': val_data['img.filename_or_obj'], 'affine': val_data['img.affine']})
         metric = metric_sum / metric_count
     return metric