remove unnecessary images from batch (#445)

Co-authored-by: Benjamin Morris <[email protected]>
AllenCellModeling · Nov 20, 2024 · 71659df · 71659df
1 parent de63972
commit 71659df
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/configs/data/im2im/segmentation_plugin.yaml b/configs/data/im2im/segmentation_plugin.yaml
@@ -76,6 +76,13 @@ transforms:
         base_image_key: ${base_image_col}
         output_name: seg
 
+      # remove keys that aren't used for training (e.g. the pre-merged targets)
+      - _target_: monai.transforms.SelectItemsd
+        keys:
+          - ${source_col}
+          - seg
+          - ${exclude_mask_col}
+
       - _target_: monai.transforms.ToTensord
         keys:
           - ${source_col}
@@ -234,6 +241,13 @@ transforms:
         base_image_key: ${base_image_col}
         output_name: seg
 
+      # remove keys that aren't used for training (e.g. the pre-merged targets)
+      - _target_: monai.transforms.SelectItemsd
+        keys:
+          - ${source_col}
+          - seg
+          - ${exclude_mask_col}
+
       - _target_: monai.transforms.ToTensord
         keys:
           - ${source_col}