add filename_or_obj metatata to multidim dataloader to create io maps

AllenCellModeling · Oct 18, 2024 · 4e6a952 · 4e6a952
1 parent 349a5e5
commit 4e6a952
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/configs/experiment/local/eval_scale1_new.yaml b/configs/experiment/local/eval_scale1_new.yaml
@@ -28,17 +28,17 @@ trainer:
 
 model:
   compile: True
-  #save_dir: /allen/aics/assay-dev/users/Suraj/EMT_Work/image_analysis_test/EMT_image_analysis/Colony_mask_training_inference/eval_whole_movie_multiscale_patch1_zarr_aws
+  save_dir: /allen/aics/assay-dev/users/Suraj/EMT_Work/image_analysis_test/EMT_image_analysis/Colony_mask_training_inference/eval_whole_movie_multiscale_patch1_zarr_aws
 
 data:
   csv_path: /allen/aics/assay-dev/users/Suraj/EMT_Work/image_analysis_test/EMT_image_analysis/Colony_mask_training_inference/sample_csv/predict_all_cells_mask_zarr_aws_v0.csv #path to csv containing test movies
   #batch_size: 1
-  _aux: 
+  _aux:
     patch_shape: [16, 128, 128]
 
 callbacks:
   predict_saving:
     _target_: cyto_dl.callbacks.ImageSaver
-    save_dir: /allen/aics/assay-dev/users/Suraj/EMT_Work/image_analysis_test/EMT_image_analysis/Colony_mask_training_inference/eval_whole_movie_multiscale_patch1_zarr_aws
+    save_dir: ${model.save_dir}
     stages: ["predict"]
-    save_input: False
+    save_input: False
diff --git a/cyto_dl/callbacks/image_saver.py b/cyto_dl/callbacks/image_saver.py
@@ -19,8 +19,9 @@ def __init__(
 
         Parameters
         ----------
-        save_dir: Union[str, Path]
-            Directory to save images
+        save_dir: Union[str, Path] 
+            Directory to save images. Only testing saves in this directory, prediction
+            saves in model save directory 
         save_every_n_epochs:int=1
             Frequency to save images
         stages:List[str]=["train", "val"]
@@ -49,6 +50,7 @@ def on_predict_batch_end(
             if outputs is None:
                 # image has already been saved
                 return
+
             for i, head_io_map in enumerate(io_map.values()):
                 for k, save_path in head_io_map.items():
                     self._save(save_path, outputs[k]["pred"][i])

diff --git a/cyto_dl/datamodules/multidim_image.py b/cyto_dl/datamodules/multidim_image.py
@@ -97,6 +97,9 @@ def _get_timepoints(self, row, img):
         step = row.get(self.time_step_column, 1)
         timepoints = range(start, stop + 1, step) if stop > 0 else range(img.dims.T)
         return list(timepoints)
+
+    def _get_filename(self, image_input_path):
+        return image_input_path.split('/')[-1].split('.')[0]
 
     def get_per_file_args(self, df):
         img_data = []
@@ -105,6 +108,8 @@ def get_per_file_args(self, df):
             img = BioImage(row[self.img_path_column])
             scenes = self._get_scenes(row, img)
             timepoints = self._get_timepoints(row, img)
+            filename = self._get_filename(row[self.img_path_column])
+
             for scene in scenes:
                 for timepoint in timepoints:
                     img_data.append(
@@ -115,6 +120,7 @@ def get_per_file_args(self, df):
                             "scene": scene,
                             "T": timepoint,
                             "original_path": row[self.img_path_column],
+                            "filename_or_obj": filename, # needs to be part of metadata to generate IO maps
                         }
                     )
         img_data.reverse()