From 83bd62091270543e77c1ae337527c3bcfda408b6 Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Wed, 20 Mar 2024 17:53:17 -0400
Subject: [PATCH 1/4] move examples

---
 dacapo/__init__.py                            |   2 +-
 .../datasplits/datasplit_generator.py         |  16 +-
 .../pipeline.py}                              |   0
 dacapo/{examples/utils.py => utils/view.py}   |   0
 {dacapo/examples => examples}/__init__.py     |   0
 examples/cosem/cosem_finetune.ipynb           | 271 ++++++++++++++++++
 examples/cosem/run_cosem.py                   | 118 ++++++++
 .../distance_task/__init__.py                 |   0
 .../distance_task/cosem_example.ipynb         |   0
 .../distance_task/cosem_example.py            |   0
 .../cosem_example_fill_in_the_blank.ipynb     |   0
 .../cosem_example_fill_in_the_blank.py        |   0
 .../cosem_finetune_example.ipynb              |   0
 .../distance_task/cosem_finetune_example.py   |   0
 .../distance_task/synthetic_example.ipynb     |   4 +-
 .../distance_task/synthetic_example.py        |   0
 {dacapo/examples => examples}/empanada.sh     |   0
 .../synthetic_source_worker.py                |   2 +-
 18 files changed, 403 insertions(+), 10 deletions(-)
 rename dacapo/{examples/random_source_pipeline.py => utils/pipeline.py} (100%)
 rename dacapo/{examples/utils.py => utils/view.py} (100%)
 rename {dacapo/examples => examples}/__init__.py (100%)
 create mode 100644 examples/cosem/cosem_finetune.ipynb
 create mode 100644 examples/cosem/run_cosem.py
 rename {dacapo/examples => examples}/distance_task/__init__.py (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_example.ipynb (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_example.py (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_example_fill_in_the_blank.ipynb (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_example_fill_in_the_blank.py (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_finetune_example.ipynb (100%)
 rename {dacapo/examples => examples}/distance_task/cosem_finetune_example.py (100%)
 rename {dacapo/examples => examples}/distance_task/synthetic_example.ipynb (99%)
 rename {dacapo/examples => examples}/distance_task/synthetic_example.py (100%)
 rename {dacapo/examples => examples}/empanada.sh (100%)
 rename {dacapo/examples => examples}/synthetic_source_worker.py (98%)

diff --git a/dacapo/__init__.py b/dacapo/__init__.py
index fe581f288..70c0c2ff1 100644
--- a/dacapo/__init__.py
+++ b/dacapo/__init__.py
@@ -1,5 +1,5 @@
 from .options import Options  # noqa
-from . import experiments  # noqa
+from . import experiments, utils  # noqa
 from .apply import apply  # noqa
 from .train import train  # noqa
 from .validate import validate  # noqa
diff --git a/dacapo/experiments/datasplits/datasplit_generator.py b/dacapo/experiments/datasplits/datasplit_generator.py
index ad998fca5..250c442b5 100644
--- a/dacapo/experiments/datasplits/datasplit_generator.py
+++ b/dacapo/experiments/datasplits/datasplit_generator.py
@@ -20,8 +20,6 @@
 
 logger = logging.getLogger(__name__)
 
-__SEPARATOR_CHARACTER = "&"
-
 
 def is_zarr_group(file_name: str, dataset: str):
     zarr_file = zarr.open(str(file_name))
@@ -187,6 +185,7 @@ def __init__(
         min_training_volume_size=8_000,  # 20**3
         raw_min=0,
         raw_max=255,
+        classes_separator_caracter = "&",
     ):
         self.name = name
         self.datasets = datasets
@@ -208,6 +207,7 @@ def __init__(
         self.min_training_volume_size = min_training_volume_size
         self.raw_min = raw_min
         self.raw_max = raw_max
+        self.classes_separator_caracter = classes_separator_caracter
 
     def __str__(self) -> str:
         return f"DataSplitGenerator:{self.name}_{self.segmentation_type}_{self.class_name}_{self.output_resolution[0]}nm"
@@ -226,7 +226,7 @@ def class_name(self, class_name):
         self._class_name = class_name
 
     def check_class_name(self, class_name):
-        datasets, classes = format_class_name(class_name)
+        datasets, classes = format_class_name(class_name,self.classes_separator_caracter)
         if self.class_name is None:
             self.class_name = classes
             if self.targets is None:
@@ -268,8 +268,12 @@ def __generate_semantic_seg_datasplit(self):
                         gt_config=gt_config,
                     )
                 )
+        if type(self.class_name) == list:
+            classes = self.classes_separator_caracter.join(self.class_name)
+        else:
+            classes = self.class_name
         return TrainValidateDataSplitConfig(
-            name=f"{self.name}_{self.segmentation_type}_{self.class_name}_{self.output_resolution[0]}nm",
+            name=f"{self.name}_{self.segmentation_type}_{classes}_{self.output_resolution[0]}nm",
             train_configs=train_dataset_configs,
             validate_configs=validation_dataset_configs,
         )
@@ -383,11 +387,11 @@ def generate_from_csv(
         )
 
 
-def format_class_name(class_name):
+def format_class_name(class_name, separator_character="&"):
     if "[" in class_name:
         if "]" not in class_name:
             raise ValueError(f"Invalid class name {class_name} missing ']'")
-        classes = class_name.split("[")[1].split("]")[0].split(__SEPARATOR_CHARACTER)
+        classes = class_name.split("[")[1].split("]")[0].split(separator_character)
         base_class_name = class_name.split("[")[0]
         return [f"{base_class_name}{c}" for c in classes], classes
     else:
diff --git a/dacapo/examples/random_source_pipeline.py b/dacapo/utils/pipeline.py
similarity index 100%
rename from dacapo/examples/random_source_pipeline.py
rename to dacapo/utils/pipeline.py
diff --git a/dacapo/examples/utils.py b/dacapo/utils/view.py
similarity index 100%
rename from dacapo/examples/utils.py
rename to dacapo/utils/view.py
diff --git a/dacapo/examples/__init__.py b/examples/__init__.py
similarity index 100%
rename from dacapo/examples/__init__.py
rename to examples/__init__.py
diff --git a/examples/cosem/cosem_finetune.ipynb b/examples/cosem/cosem_finetune.ipynb
new file mode 100644
index 000000000..342bd16d5
--- /dev/null
+++ b/examples/cosem/cosem_finetune.ipynb
@@ -0,0 +1,271 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating FileConfigStore:\n",
+      "\tpath: /groups/scicompsoft/home/zouinkhim/dacapo/configs\n"
+     ]
+    }
+   ],
+   "source": [
+    "from dacapo.store.create_store import create_config_store\n",
+    "\n",
+    "config_store = create_config_store()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# config_store.retrieve_datasplit_config_names()\n",
+    "# config_store.retrieve_task_config_names()\n",
+    "# config_store.retrieve_architecture_config_names()\n",
+    "# config_store.retrieve_trainer_config_names()\n",
+    "# config_store.retrieve_run_config_names()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datasplit_config = config_store.retrieve_datasplit_config(\"cosem_example_semantic_mito_4nm\")\n",
+    "task_config = config_store.retrieve_task_config(\"cosem_distance_task_4nm\")\n",
+    "architecture_config = config_store.retrieve_architecture_config(\"upsample_unet\")\n",
+    "trainer_config = config_store.retrieve_trainer_config(\"cosem_finetune\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Neuroglancer link: http://h10u12.int.janelia.org:36333/v/025325da9eea3bd4848c24ab6424ef528d4ea174/\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "datasplit = datasplit_config.datasplit_type(datasplit_config)\n",
+    "viewer = datasplit._neuroglancer()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dacapo.experiments.starts import CosemStartConfig\n",
+    "start_config = CosemStartConfig(\"setup04\", \"1820500\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dacapo.experiments import RunConfig\n",
+    "\n",
+    "run_config = RunConfig(\n",
+    "        name=\"cosem_distance_run_4nm_finetune\",\n",
+    "        datasplit_config=datasplit_config,\n",
+    "        task_config=task_config,\n",
+    "        architecture_config=architecture_config,\n",
+    "        trainer_config=trainer_config,\n",
+    "        num_iterations=2000,\n",
+    "        validation_interval=500,\n",
+    "        repetition=0,\n",
+    "        start_config=start_config,\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# visualize "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/cosem_start.py:Starter model resolution: input [8 8 8] output [4 4 4], Make sure to set the correct resolution for the input data.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:loading weights from run setup04, criterion: 1820500, old_head ['ecs', 'pm', 'mito', 'mito_mem', 'ves', 'ves_mem', 'endo', 'endo_mem', 'er', 'er_mem', 'eres', 'nuc', 'mt', 'mt_out'], new_head: ['mito']\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matching heads from run setup04, criterion: 1820500\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:old head: ['ecs', 'pm', 'mito', 'mito_mem', 'ves', 'ves_mem', 'endo', 'endo_mem', 'er', 'er_mem', 'eres', 'nuc', 'mt', 'mt_out']\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:new head: ['mito']\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Not loading weights for setup04.\n",
+      "Creating local weights store in directory /groups/scicompsoft/home/zouinkhim/dacapo\n",
+      "Creating local weights store in directory /groups/scicompsoft/home/zouinkhim/dacapo\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:Unable to load model in strict mode. Loading flexibly.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matching head for mito.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matched head for mito.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from dacapo.experiments.run import Run\n",
+    "run = Run(run_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Predicting with input size (4608, 4608, 4608), output size (1696, 1696, 1696)\n",
+      "Total input ROI: [4544:11456, 544:5456, 4544:9456] (6912, 4912, 4912), output ROI: [6000:10000, 2000:4000, 6000:8000] (4000, 2000, 2000)\n",
+      "Running blockwise prediction with worker_file:  /groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py\n",
+      "Defining worker with command:  ['/groups/scicompsoft/home/zouinkhim/miniconda3/envs/dacapo_11/bin/python', '/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py', 'start-worker', '--run-name', 'cosem_distance_run_4nm_finetune', '--input_container', '/misc/public/dacapo_learnathon/jrc_hela-2.zarr', '--input_dataset', 'recon-1/em/fibsem-uint8/s2', '--output_container', '/nrs/cellmap/zouinkhim/predictions/test_predict/hela2.zarr', '--output_dataset', 'prediction_cosem_distance_run_4nm_finetune_None']\n",
+      "Running blockwise with worker_file:  /groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py\n",
+      "Using compute context: LocalTorch(_device=None, oom_limit=4.2)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/json": {
+       "ascii": false,
+       "bar_format": null,
+       "colour": null,
+       "elapsed": 0.0072765350341796875,
+       "initial": 0,
+       "n": 0,
+       "ncols": null,
+       "nrows": null,
+       "postfix": null,
+       "prefix": "predict_worker2024-03-20_17-35-30 ▶",
+       "rate": null,
+       "total": 2,
+       "unit": "blocks",
+       "unit_divisor": 1000,
+       "unit_scale": false
+      },
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a15ad763962d47e2a356fa984c85fc0c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "predict_worker2024-03-20_17-35-30 ▶:   0%|          | 0/2 [00:00<?, ?blocks/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Execution Summary\n",
+      "-----------------\n",
+      "\n",
+      "  Task predict_worker2024-03-20_17-35-30:\n",
+      "\n",
+      "    num blocks : 2\n",
+      "    completed ✔: 2 (skipped 0)\n",
+      "    failed    ✗: 0\n",
+      "    orphaned  ∅: 0\n",
+      "\n",
+      "    all blocks processed successfully\n",
+      "Done predicting.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from dacapo.predict import predict\n",
+    "input_container = \"/misc/public/dacapo_learnathon/jrc_hela-2.zarr\"\n",
+    "input_dataset = \"recon-1/em/fibsem-uint8/s2\"\n",
+    "test_out_path = \"/nrs/cellmap/zouinkhim/predictions/test_predict/hela2.zarr\"\n",
+    "out_roi = \"[6000:10000,2000:4000,6000:8000]\"\n",
+    "\n",
+    "predict(\n",
+    "    \"cosem_distance_run_4nm_finetune\",\n",
+    "    None,\n",
+    "    input_container,\n",
+    "    input_dataset,\n",
+    "    test_out_path,\n",
+    "    num_workers=1,\n",
+    "    overwrite=False,\n",
+    "    output_dtype=\"float32\",\n",
+    "    output_roi=out_roi,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dacapo_11",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/cosem/run_cosem.py b/examples/cosem/run_cosem.py
new file mode 100644
index 000000000..599554b0f
--- /dev/null
+++ b/examples/cosem/run_cosem.py
@@ -0,0 +1,118 @@
+# %%
+import dacapo
+# %%
+from dacapo.store.create_store import create_config_store
+
+config_store = create_config_store()
+# %%
+from dacapo.experiments.datasplits import DataSplitGenerator
+from funlib.geometry import Coordinate
+
+input_resolution = Coordinate(8, 8, 8)
+output_resolution = Coordinate(4, 4, 4)
+datasplit_config = DataSplitGenerator.generate_from_csv(
+    "/misc/public/dacapo_learnathon/datasplit_csvs/cosem_example.csv",
+    input_resolution,
+    output_resolution,
+).compute()
+
+datasplit = datasplit_config.datasplit_type(datasplit_config)
+viewer = datasplit._neuroglancer()
+config_store.store_datasplit_config(datasplit_config)
+# %%
+from dacapo.experiments.tasks import DistanceTaskConfig
+
+task_config = DistanceTaskConfig(
+    name="cosem_distance_task_4nm",
+    channels=["mito"],
+    clip_distance=40.0,
+    tol_distance=40.0,
+    scale_factor=80.0,
+)
+config_store.store_task_config(task_config)
+# %%
+from dacapo.experiments.architectures import CNNectomeUNetConfig
+
+architecture_config = CNNectomeUNetConfig(
+    name="upsample_unet",
+    input_shape=Coordinate(216, 216, 216),
+    eval_shape_increase=Coordinate(72, 72, 72),
+    fmaps_in=1,
+    num_fmaps=12,
+    fmaps_out=72,
+    fmap_inc_factor=6,
+    downsample_factors=[(2, 2, 2), (3, 3, 3), (3, 3, 3)],
+    constant_upsample=True,
+    upsample_factors=[(2, 2, 2)],
+)
+config_store.store_architecture_config(architecture_config)
+# %%
+from dacapo.experiments.trainers.gp_augments import (
+    ElasticAugmentConfig,
+    GammaAugmentConfig,
+    IntensityAugmentConfig,
+    IntensityScaleShiftAugmentConfig,
+)
+from dacapo.experiments.trainers import GunpowderTrainerConfig
+trainer_config = GunpowderTrainerConfig(
+    name="cosem_finetune",
+    batch_size=1,
+    learning_rate=0.0001,
+    num_data_fetchers=20,
+    augments=[
+        ElasticAugmentConfig(
+            control_point_spacing=[100, 100, 100],
+            control_point_displacement_sigma=[10.0, 10.0, 10.0],
+            rotation_interval=(0.0, 1.5707963267948966),
+            subsample=8,
+            uniform_3d_rotation=True,
+        ),
+        IntensityAugmentConfig(scale=(0.25, 1.75), shift=(-0.5, 0.35), clip=True),
+        GammaAugmentConfig(gamma_range=(0.5, 2.0)),
+        IntensityScaleShiftAugmentConfig(scale=2.0, shift=-1.0),
+    ],
+    snapshot_interval=10000,
+    min_masked=0.05,
+    clip_raw=True,
+)
+config_store.store_trainer_config(trainer_config)
+# %%
+from dacapo.experiments import RunConfig
+from dacapo.experiments.run import Run
+
+from dacapo.experiments.starts import CosemStartConfig
+
+# We will now download a pretrained cosem model and finetune from that model. It will only have to download the first time it is used.
+
+start_config = CosemStartConfig("setup04", "1820500")
+# %%
+iterations = 2000
+validation_interval = iterations // 2
+repetitions = 1
+for i in range(repetitions):
+    run_config = RunConfig(
+        name="cosem_distance_run_4nm_finetune",
+        datasplit_config=datasplit_config,
+        task_config=task_config,
+        architecture_config=architecture_config,
+        trainer_config=trainer_config,
+        num_iterations=iterations,
+        validation_interval=validation_interval,
+        repetition=i,
+        start_config=start_config,
+    )
+
+    print(run_config.name)
+    config_store.store_run_config(run_config)
+# %%
+from dacapo.train import train_run
+from dacapo.experiments.run import Run
+from dacapo.store.create_store import create_config_store
+
+config_store = create_config_store()
+
+run = Run(config_store.retrieve_run_config("cosem_distance_run_4nm_finetune"))
+# train_run(run)
+# %%
+
+# %%
diff --git a/dacapo/examples/distance_task/__init__.py b/examples/distance_task/__init__.py
similarity index 100%
rename from dacapo/examples/distance_task/__init__.py
rename to examples/distance_task/__init__.py
diff --git a/dacapo/examples/distance_task/cosem_example.ipynb b/examples/distance_task/cosem_example.ipynb
similarity index 100%
rename from dacapo/examples/distance_task/cosem_example.ipynb
rename to examples/distance_task/cosem_example.ipynb
diff --git a/dacapo/examples/distance_task/cosem_example.py b/examples/distance_task/cosem_example.py
similarity index 100%
rename from dacapo/examples/distance_task/cosem_example.py
rename to examples/distance_task/cosem_example.py
diff --git a/dacapo/examples/distance_task/cosem_example_fill_in_the_blank.ipynb b/examples/distance_task/cosem_example_fill_in_the_blank.ipynb
similarity index 100%
rename from dacapo/examples/distance_task/cosem_example_fill_in_the_blank.ipynb
rename to examples/distance_task/cosem_example_fill_in_the_blank.ipynb
diff --git a/dacapo/examples/distance_task/cosem_example_fill_in_the_blank.py b/examples/distance_task/cosem_example_fill_in_the_blank.py
similarity index 100%
rename from dacapo/examples/distance_task/cosem_example_fill_in_the_blank.py
rename to examples/distance_task/cosem_example_fill_in_the_blank.py
diff --git a/dacapo/examples/distance_task/cosem_finetune_example.ipynb b/examples/distance_task/cosem_finetune_example.ipynb
similarity index 100%
rename from dacapo/examples/distance_task/cosem_finetune_example.ipynb
rename to examples/distance_task/cosem_finetune_example.ipynb
diff --git a/dacapo/examples/distance_task/cosem_finetune_example.py b/examples/distance_task/cosem_finetune_example.py
similarity index 100%
rename from dacapo/examples/distance_task/cosem_finetune_example.py
rename to examples/distance_task/cosem_finetune_example.py
diff --git a/dacapo/examples/distance_task/synthetic_example.ipynb b/examples/distance_task/synthetic_example.ipynb
similarity index 99%
rename from dacapo/examples/distance_task/synthetic_example.ipynb
rename to examples/distance_task/synthetic_example.ipynb
index 986cc68f1..fa9bad3a5 100644
--- a/dacapo/examples/distance_task/synthetic_example.ipynb
+++ b/examples/distance_task/synthetic_example.ipynb
@@ -123,8 +123,8 @@
     "\n",
     "from pathlib import Path\n",
     "from dacapo import Options\n",
-    "from dacapo.examples.utils import get_viewer\n",
-    "from dacapo.examples.synthetic_source_worker import generate_synthetic_dataset\n",
+    "from dacapo.utils.view import get_viewer\n",
+    "from ..synthetic_source_worker import generate_synthetic_dataset\n",
     "from funlib.geometry import Coordinate\n",
     "from funlib.persistence import open_ds\n",
     "\n",
diff --git a/dacapo/examples/distance_task/synthetic_example.py b/examples/distance_task/synthetic_example.py
similarity index 100%
rename from dacapo/examples/distance_task/synthetic_example.py
rename to examples/distance_task/synthetic_example.py
diff --git a/dacapo/examples/empanada.sh b/examples/empanada.sh
similarity index 100%
rename from dacapo/examples/empanada.sh
rename to examples/empanada.sh
diff --git a/dacapo/examples/synthetic_source_worker.py b/examples/synthetic_source_worker.py
similarity index 98%
rename from dacapo/examples/synthetic_source_worker.py
rename to examples/synthetic_source_worker.py
index fb394cc36..cd2744fa0 100644
--- a/dacapo/examples/synthetic_source_worker.py
+++ b/examples/synthetic_source_worker.py
@@ -1,4 +1,4 @@
-from dacapo.examples.random_source_pipeline import random_source_pipeline
+from dacapo.utils.pipeline import random_source_pipeline
 import gunpowder as gp
 
 

From 18859c5b5aaffda5f6192f7cf8606e2266806fae Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Wed, 20 Mar 2024 22:58:59 -0400
Subject: [PATCH 2/4] hela example

---
 examples/cosem/cosem_finetune.ipynb     | 18 ++++++++++--------
 examples/cosem/submit_train.sh          |  1 +
 examples/cosem/train_hela.py            | 10 ++++++++++
 examples/distance_task/cosem_example.py |  2 +-
 4 files changed, 22 insertions(+), 9 deletions(-)
 create mode 100644 examples/cosem/submit_train.sh
 create mode 100644 examples/cosem/train_hela.py

diff --git a/examples/cosem/cosem_finetune.ipynb b/examples/cosem/cosem_finetune.ipynb
index 342bd16d5..e97d60710 100644
--- a/examples/cosem/cosem_finetune.ipynb
+++ b/examples/cosem/cosem_finetune.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -47,14 +47,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Neuroglancer link: http://h10u12.int.janelia.org:36333/v/025325da9eea3bd4848c24ab6424ef528d4ea174/\n"
+      "Neuroglancer link: http://h10u12.int.janelia.org:21785/v/f2cebdf760ce1844f43003a61d3ac023e59e6361/\n"
      ]
     }
    ],
@@ -66,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -76,7 +76,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -92,7 +92,9 @@
     "        validation_interval=500,\n",
     "        repetition=0,\n",
     "        start_config=start_config,\n",
-    "    )"
+    "    )\n",
+    "# config_store.delete_run_config(run_config.name)\n",
+    "config_store.store_run_config(run_config)"
    ]
   },
   {
diff --git a/examples/cosem/submit_train.sh b/examples/cosem/submit_train.sh
new file mode 100644
index 000000000..4a7ec970f
--- /dev/null
+++ b/examples/cosem/submit_train.sh
@@ -0,0 +1 @@
+bsub -J Jupyter -q gpu_tesla -gpu "num=1" -n16 -We 120 python train_hela.py 
\ No newline at end of file
diff --git a/examples/cosem/train_hela.py b/examples/cosem/train_hela.py
new file mode 100644
index 000000000..33cc37705
--- /dev/null
+++ b/examples/cosem/train_hela.py
@@ -0,0 +1,10 @@
+import dacapo
+from dacapo.store.create_store import create_config_store
+
+config_store = create_config_store()
+run_config = config_store.retrieve_run_config("cosem_distance_run_4nm_finetune")
+
+from dacapo.experiments.run import Run
+run = Run(run_config)
+from dacapo.train import train_run
+train_run(run)
diff --git a/examples/distance_task/cosem_example.py b/examples/distance_task/cosem_example.py
index da07f091c..9e990f0e6 100644
--- a/examples/distance_task/cosem_example.py
+++ b/examples/distance_task/cosem_example.py
@@ -185,7 +185,7 @@
 repetitions = 1
 for i in range(repetitions):
     run_config = RunConfig(
-        name="cosem_distance_run_4nm",
+        name="scratch_cosem_distance_run_4nm",
         # # NOTE: This is a template for the name of the run. You can customize it as you see fit.
         # name=("_").join(
         #     [

From cf18ab93cda75ea5a607c89d00bf2cb4792668e1 Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Thu, 21 Mar 2024 00:33:29 -0400
Subject: [PATCH 3/4] cosem example

---
 examples/cosem/cosem_finetune.ipynb | 151 +++++++++++++++++++++-------
 examples/cosem/run_cosem.py         |  11 +-
 examples/cosem/submit_train.sh      |   2 +-
 examples/cosem/train_hela.py        |   2 +-
 4 files changed, 123 insertions(+), 43 deletions(-)
 mode change 100644 => 100755 examples/cosem/submit_train.sh

diff --git a/examples/cosem/cosem_finetune.ipynb b/examples/cosem/cosem_finetune.ipynb
index e97d60710..24d33dba0 100644
--- a/examples/cosem/cosem_finetune.ipynb
+++ b/examples/cosem/cosem_finetune.ipynb
@@ -35,26 +35,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
     "datasplit_config = config_store.retrieve_datasplit_config(\"cosem_example_semantic_mito_4nm\")\n",
     "task_config = config_store.retrieve_task_config(\"cosem_distance_task_4nm\")\n",
     "architecture_config = config_store.retrieve_architecture_config(\"upsample_unet\")\n",
-    "trainer_config = config_store.retrieve_trainer_config(\"cosem_finetune\")"
+    "trainer_config = config_store.retrieve_trainer_config(\"cosem_finetune2\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Neuroglancer link: http://h10u12.int.janelia.org:21785/v/f2cebdf760ce1844f43003a61d3ac023e59e6361/\n"
+      "Neuroglancer link: http://h10u24.int.janelia.org:36025/v/08c3c01bb86d2c5a555fb96ccc8b87cf1850d461/\n"
      ]
     }
    ],
@@ -66,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -76,14 +76,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
     "from dacapo.experiments import RunConfig\n",
     "\n",
     "run_config = RunConfig(\n",
-    "        name=\"cosem_distance_run_4nm_finetune\",\n",
+    "        name=\"cosem_distance_run_4nm_finetune3\",\n",
     "        datasplit_config=datasplit_config,\n",
     "        task_config=task_config,\n",
     "        architecture_config=architecture_config,\n",
@@ -99,25 +99,57 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 30,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/cosem_start.py:Starter model resolution: input [8 8 8] output [4 4 4], Make sure to set the correct resolution for the input data.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:loading weights from run setup04, criterion: 1820500, old_head ['ecs', 'pm', 'mito', 'mito_mem', 'ves', 'ves_mem', 'endo', 'endo_mem', 'er', 'er_mem', 'eres', 'nuc', 'mt', 'mt_out'], new_head: ['mito']\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matching heads from run setup04, criterion: 1820500\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:old head: ['ecs', 'pm', 'mito', 'mito_mem', 'ves', 'ves_mem', 'endo', 'endo_mem', 'er', 'er_mem', 'eres', 'nuc', 'mt', 'mt_out']\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:new head: ['mito']\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Not loading weights for setup04.\n",
+      "Creating local weights store in directory /groups/scicompsoft/home/zouinkhim/dacapo\n",
+      "Creating local weights store in directory /groups/scicompsoft/home/zouinkhim/dacapo\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:Unable to load model in strict mode. Loading flexibly.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matching head for mito.\n",
+      "WARNING:/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/experiments/starts/start.py:matched head for mito.\n"
+     ]
+    }
+   ],
    "source": [
-    "# visualize "
+    "from dacapo.experiments.run import Run\n",
+    "run = Run(run_config)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating FileConfigStore:\n",
+      "\tpath: /groups/scicompsoft/home/zouinkhim/dacapo/configs\n"
+     ]
+    },
     {
      "name": "stderr",
      "output_type": "stream",
@@ -149,23 +181,31 @@
     }
    ],
    "source": [
+    "\n",
+    "from dacapo.train import train_run\n",
     "from dacapo.experiments.run import Run\n",
-    "run = Run(run_config)"
+    "from dacapo.store.create_store import create_config_store\n",
+    "\n",
+    "config_store = create_config_store()\n",
+    "\n",
+    "run = Run(config_store.retrieve_run_config(\"cosem_distance_run_4nm_finetune2\"))\n",
+    "# # we already trained it, so we will just load the weights\n",
+    "# # train_run(run)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Predicting with input size (4608, 4608, 4608), output size (1696, 1696, 1696)\n",
-      "Total input ROI: [4544:11456, 544:5456, 4544:9456] (6912, 4912, 4912), output ROI: [6000:10000, 2000:4000, 6000:8000] (4000, 2000, 2000)\n",
+      "Predicting with input size (2304, 2304, 2304), output size (848, 848, 848)\n",
+      "Total input ROI: [3672:7128, 472:3928, 28072:31528] (3456, 3456, 3456), output ROI: [4400:6400, 1200:3200, 28800:30800] (2000, 2000, 2000)\n",
       "Running blockwise prediction with worker_file:  /groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py\n",
-      "Defining worker with command:  ['/groups/scicompsoft/home/zouinkhim/miniconda3/envs/dacapo_11/bin/python', '/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py', 'start-worker', '--run-name', 'cosem_distance_run_4nm_finetune', '--input_container', '/misc/public/dacapo_learnathon/jrc_hela-2.zarr', '--input_dataset', 'recon-1/em/fibsem-uint8/s2', '--output_container', '/nrs/cellmap/zouinkhim/predictions/test_predict/hela2.zarr', '--output_dataset', 'prediction_cosem_distance_run_4nm_finetune_None']\n",
+      "Defining worker with command:  ['/groups/scicompsoft/home/zouinkhim/miniconda3/envs/dacapo_11/bin/python', '/groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py', 'start-worker', '--run-name', 'cosem_distance_run_4nm_finetune2', '--input_container', '/misc/public/dacapo_learnathon/jrc_hela-2.zarr', '--input_dataset', 'recon-1/em/fibsem-uint8/s1', '--output_container', '/nrs/cellmap/zouinkhim/predictions/test_predict/hela2__v2_0_s1.zarr', '--output_dataset', 'prediction_cosem_distance_run_4nm_finetune2_None']\n",
       "Running blockwise with worker_file:  /groups/cellmap/cellmap/zouinkhim/dacapo_release/dacapo/dacapo/blockwise/predict_worker.py\n",
       "Using compute context: LocalTorch(_device=None, oom_limit=4.2)\n"
      ]
@@ -176,26 +216,26 @@
        "ascii": false,
        "bar_format": null,
        "colour": null,
-       "elapsed": 0.0072765350341796875,
+       "elapsed": 0.007761716842651367,
        "initial": 0,
        "n": 0,
        "ncols": null,
        "nrows": null,
        "postfix": null,
-       "prefix": "predict_worker2024-03-20_17-35-30 ▶",
+       "prefix": "predict_worker2024-03-21_00-09-32 ▶",
        "rate": null,
-       "total": 2,
+       "total": 8,
        "unit": "blocks",
        "unit_divisor": 1000,
        "unit_scale": false
       },
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a15ad763962d47e2a356fa984c85fc0c",
+       "model_id": "d323afb243a6445f99741f3523599802",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "predict_worker2024-03-20_17-35-30 ▶:   0%|          | 0/2 [00:00<?, ?blocks/s]"
+       "predict_worker2024-03-21_00-09-32 ▶:   0%|          | 0/8 [00:00<?, ?blocks/s]"
       ]
      },
      "metadata": {},
@@ -209,10 +249,10 @@
       "Execution Summary\n",
       "-----------------\n",
       "\n",
-      "  Task predict_worker2024-03-20_17-35-30:\n",
+      "  Task predict_worker2024-03-21_00-09-32:\n",
       "\n",
-      "    num blocks : 2\n",
-      "    completed ✔: 2 (skipped 0)\n",
+      "    num blocks : 8\n",
+      "    completed ✔: 8 (skipped 0)\n",
       "    failed    ✗: 0\n",
       "    orphaned  ∅: 0\n",
       "\n",
@@ -224,12 +264,13 @@
    "source": [
     "from dacapo.predict import predict\n",
     "input_container = \"/misc/public/dacapo_learnathon/jrc_hela-2.zarr\"\n",
-    "input_dataset = \"recon-1/em/fibsem-uint8/s2\"\n",
-    "test_out_path = \"/nrs/cellmap/zouinkhim/predictions/test_predict/hela2.zarr\"\n",
-    "out_roi = \"[6000:10000,2000:4000,6000:8000]\"\n",
+    "input_dataset = \"recon-1/em/fibsem-uint8/s1\"\n",
+    "test_out_path = \"/nrs/cellmap/zouinkhim/predictions/test_predict/hela2__v2_0_s1.zarr\"\n",
+    "# out_roi = \"[6000:10000,2000:4000,6000:8000]\"\n",
+    "out_roi = \"[4400:6400,1200:3200,28800:30800]\"\n",
     "\n",
     "predict(\n",
-    "    \"cosem_distance_run_4nm_finetune\",\n",
+    "    run,\n",
     "    None,\n",
     "    input_container,\n",
     "    input_dataset,\n",
@@ -241,12 +282,50 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# visualize in neuroglancer"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "# from dacapo.train import train_run\n",
+    "# train_run(run)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dacapo.predict import predict\n",
+    "input_container = \"/misc/public/dacapo_learnathon/jrc_hela-2.zarr\"\n",
+    "input_dataset = \"recon-1/em/fibsem-uint8/s1\"\n",
+    "test_out_path = \"/nrs/cellmap/zouinkhim/predictions/test_predict/hela2__v2_2000_s1.zarr\"\n",
+    "# out_roi = \"[6000:10000,2000:4000,6000:8000]\"\n",
+    "out_roi = \"[4400:6400,1200:3200,28800:30800]\"\n",
+    "\n",
+    "predict(\n",
+    "    run,\n",
+    "    2000,\n",
+    "    input_container,\n",
+    "    input_dataset,\n",
+    "    test_out_path,\n",
+    "    num_workers=1,\n",
+    "    overwrite=False,\n",
+    "    output_dtype=\"float32\",\n",
+    "    output_roi=out_roi,\n",
+    ")"
+   ]
   }
  ],
  "metadata": {
diff --git a/examples/cosem/run_cosem.py b/examples/cosem/run_cosem.py
index 599554b0f..2b2f73aab 100644
--- a/examples/cosem/run_cosem.py
+++ b/examples/cosem/run_cosem.py
@@ -55,9 +55,9 @@
 )
 from dacapo.experiments.trainers import GunpowderTrainerConfig
 trainer_config = GunpowderTrainerConfig(
-    name="cosem_finetune",
+    name="cosem_finetune2",
     batch_size=1,
-    learning_rate=0.0001,
+    learning_rate=0.001,
     num_data_fetchers=20,
     augments=[
         ElasticAugmentConfig(
@@ -72,8 +72,8 @@
         IntensityScaleShiftAugmentConfig(scale=2.0, shift=-1.0),
     ],
     snapshot_interval=10000,
-    min_masked=0.05,
-    clip_raw=True,
+    # min_masked=0.05,
+    clip_raw=False,
 )
 config_store.store_trainer_config(trainer_config)
 # %%
@@ -91,7 +91,7 @@
 repetitions = 1
 for i in range(repetitions):
     run_config = RunConfig(
-        name="cosem_distance_run_4nm_finetune",
+        name="cosem_distance_run_4nm_finetune2",
         datasplit_config=datasplit_config,
         task_config=task_config,
         architecture_config=architecture_config,
@@ -112,6 +112,7 @@
 config_store = create_config_store()
 
 run = Run(config_store.retrieve_run_config("cosem_distance_run_4nm_finetune"))
+# we already trained it, so we will just load the weights
 # train_run(run)
 # %%
 
diff --git a/examples/cosem/submit_train.sh b/examples/cosem/submit_train.sh
old mode 100644
new mode 100755
index 4a7ec970f..6526e9b92
--- a/examples/cosem/submit_train.sh
+++ b/examples/cosem/submit_train.sh
@@ -1 +1 @@
-bsub -J Jupyter -q gpu_tesla -gpu "num=1" -n16 -We 120 python train_hela.py 
\ No newline at end of file
+bsub -P cellmap -J train_cosem -q gpu_tesla -gpu "num=1" -n16 -We 120 python train_hela.py 
\ No newline at end of file
diff --git a/examples/cosem/train_hela.py b/examples/cosem/train_hela.py
index 33cc37705..21521aeee 100644
--- a/examples/cosem/train_hela.py
+++ b/examples/cosem/train_hela.py
@@ -2,7 +2,7 @@
 from dacapo.store.create_store import create_config_store
 
 config_store = create_config_store()
-run_config = config_store.retrieve_run_config("cosem_distance_run_4nm_finetune")
+run_config = config_store.retrieve_run_config("cosem_distance_run_4nm_finetune3")
 
 from dacapo.experiments.run import Run
 run = Run(run_config)

From e223a284532c768f9191b0279ad972629dbabf40 Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Thu, 21 Mar 2024 00:41:10 -0400
Subject: [PATCH 4/4] black format

---
 dacapo/experiments/datasplits/datasplit_generator.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/dacapo/experiments/datasplits/datasplit_generator.py b/dacapo/experiments/datasplits/datasplit_generator.py
index 250c442b5..8f177e187 100644
--- a/dacapo/experiments/datasplits/datasplit_generator.py
+++ b/dacapo/experiments/datasplits/datasplit_generator.py
@@ -185,7 +185,7 @@ def __init__(
         min_training_volume_size=8_000,  # 20**3
         raw_min=0,
         raw_max=255,
-        classes_separator_caracter = "&",
+        classes_separator_caracter="&",
     ):
         self.name = name
         self.datasets = datasets
@@ -226,7 +226,9 @@ def class_name(self, class_name):
         self._class_name = class_name
 
     def check_class_name(self, class_name):
-        datasets, classes = format_class_name(class_name,self.classes_separator_caracter)
+        datasets, classes = format_class_name(
+            class_name, self.classes_separator_caracter
+        )
         if self.class_name is None:
             self.class_name = classes
             if self.targets is None: