From 349a5e5c33df89c114159aa31474cbe30013b580 Mon Sep 17 00:00:00 2001
From: Suraj Mishra <suraj.mishra@alleninstitute.org>
Date: Wed, 16 Oct 2024 10:28:00 -0700
Subject: [PATCH] Trying to use the Zarr loader

---
 ...segmentation_all_cells_mask_from_zarr.yaml | 18 +++++++
 configs/experiment/local/eval_scale1_new.yaml | 44 +++++++++++++++
 .../local/segmentation_all_cells_mask.yaml    | 54 +++++++++++++++++++
 3 files changed, 116 insertions(+)
 create mode 100644 configs/data/local/segmentation_all_cells_mask_from_zarr.yaml
 create mode 100644 configs/experiment/local/eval_scale1_new.yaml
 create mode 100644 configs/model/local/segmentation_all_cells_mask.yaml

diff --git a/configs/data/local/segmentation_all_cells_mask_from_zarr.yaml b/configs/data/local/segmentation_all_cells_mask_from_zarr.yaml
new file mode 100644
index 000000000..e34a77246
--- /dev/null
+++ b/configs/data/local/segmentation_all_cells_mask_from_zarr.yaml
@@ -0,0 +1,18 @@
+_target_: cyto_dl.datamodules.multidim_image.make_multidim_image_dataloader #MultiDimImageDataset #dataframe.DataframeDatamodule
+
+num_workers: 0
+batch_size: 1
+pin_memory: True
+#persistent_workers: False
+
+
+csv_path:
+img_path_column: movie_path
+channel_column: bf_channel
+#fms_id_column: fms_id
+out_key: ${source_col}
+transforms:
+  - _target_: monai.transforms.ToTensor
+  - _target_: monai.transforms.NormalizeIntensity
+
+
diff --git a/configs/experiment/local/eval_scale1_new.yaml b/configs/experiment/local/eval_scale1_new.yaml
new file mode 100644
index 000000000..552cd8243
--- /dev/null
+++ b/configs/experiment/local/eval_scale1_new.yaml
@@ -0,0 +1,44 @@
+# @package _global_
+# to execute this experiment run:
+# python train.py experiment=example
+defaults:
+  - override /data: local/segmentation_all_cells_mask_from_zarr.yaml
+  - override /model: local/segmentation_all_cells_mask.yaml
+  - override /callbacks: default.yaml
+  - override /trainer: gpu.yaml
+  - override /logger: mlflow.yaml
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+tags: ["dev"]
+seed: 12345
+
+experiment_name: train_cytogfp_seg_model_v0
+run_name: test_v0_patch1
+source_col: raw
+target_col: seg
+spatial_dims: 3
+raw_im_channels: 1
+checkpoint:
+  ckpt_path: /allen/aics/assay-dev/users/Suraj/EMT_Work/image_analysis_test/EMT_image_analysis/Colony_mask_training_inference/all_cells_mask_seg_model_checkpoint.ckpt
+trainer:
+  max_epochs: 1000
+  devices: [0]
+
+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
+
+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: 
+    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
+    stages: ["predict"]
+    save_input: False
\ No newline at end of file
diff --git a/configs/model/local/segmentation_all_cells_mask.yaml b/configs/model/local/segmentation_all_cells_mask.yaml
new file mode 100644
index 000000000..c9994ac02
--- /dev/null
+++ b/configs/model/local/segmentation_all_cells_mask.yaml
@@ -0,0 +1,54 @@
+_target_: cyto_dl.models.im2im.MultiTaskIm2Im
+
+save_images_every_n_epochs: 50
+save_dir: ${paths.output_dir}
+
+x_key: ${source_col}
+# patch_shape:
+
+backbone:
+  _target_: monai.networks.nets.DynUNet
+  spatial_dims: ${spatial_dims}
+  in_channels: ${raw_im_channels}
+  out_channels: 1
+  strides: [1, 2, 2, 2]
+  kernel_size: [3, 3, 3, 3]
+  upsample_kernel_size: [2, 2, 2]
+  dropout: 0.0
+  res_block: True
+
+task_heads: ${kv_to_dict:${model._aux._tasks}}
+
+optimizer:
+  generator:
+    _partial_: True
+    _target_: torch.optim.Adam
+    lr: 0.0001
+    weight_decay: 0.0001
+
+lr_scheduler:
+  generator:
+    _partial_: True
+    _target_: torch.optim.lr_scheduler.ExponentialLR
+    gamma: 0.995
+
+inference_args:
+  sw_batch_size: 1
+  roi_size: ${data._aux.patch_shape}
+
+_aux:
+  _tasks:
+    - - ${target_col}
+      - _target_: cyto_dl.nn.BaseHead
+        loss:
+          _target_: monai.losses.GeneralizedDiceFocalLoss  ##Main loss
+          sigmoid: True
+        postprocess:
+          input:
+            _target_: cyto_dl.models.im2im.utils.postprocessing.ActThreshLabel
+            rescale_dtype: numpy.uint8
+          prediction:
+            _target_: cyto_dl.models.im2im.utils.postprocessing.ActThreshLabel
+            activation:
+              _target_: torch.nn.Sigmoid
+            rescale_dtype: numpy.uint8