From 71659df96926997a5fb6fa0873b5c0b9b3f4c268 Mon Sep 17 00:00:00 2001
From: benjijamorris <54606172+benjijamorris@users.noreply.github.com>
Date: Wed, 20 Nov 2024 09:19:38 -0800
Subject: [PATCH] remove unnecessary images from batch (#445)

Co-authored-by: Benjamin Morris <benjamin.morris@prd-aics-dcp-038.corp.alleninstitute.org>
---
 configs/data/im2im/segmentation_plugin.yaml | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/configs/data/im2im/segmentation_plugin.yaml b/configs/data/im2im/segmentation_plugin.yaml
index 26bbb7ef..9f3aaf59 100644
--- 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}