Merge pull request #51 from janelia-cellmap/docs

Tests + fix bugs + scriptable

Author: mzouink

cellmap_flow/blockwise/blockwise_processor.py

@@ -16,24 +16,12 @@
 from funlib.persistence import prepare_ds, open_ds, Array
 from pathlib import Path
 
-# from cellmap_flow.globals import Flow
-import cellmap_flow.globals as g
+from cellmap_flow.globals import g
 from cellmap_flow.utils.web_utils import encode_to_str, decode_to_json
 
 logger = logging.getLogger(__name__)
 
 
-def get_output_dtype(model_output):
-    p_dtype = model_output
-    # g = Flow()
-    if len(g.postprocess) > 0:
-        for postprocess in g.postprocess[::-1]:
-            if postprocess.dtype:
-                p_dtype = postprocess.dtype
-                break
-    return p_dtype
-
-
 def get_process_dataset(json_data: str):
     logger.error(f"json data: {json_data}")
     input_norm_fns = get_normalizations(json_data[INPUT_NORM_DICT_KEY])
@@ -60,7 +48,6 @@ def __init__(self, yaml_config: str, create=True):
         self.output_path = self.config["output_path"]
         self.output_path = Path(self.output_path)
 
-
         output_channels = None
         if "output_channels" in self.config:
             output_channels = self.config["output_channels"].split(",")
@@ -112,7 +99,7 @@ def __init__(self, yaml_config: str, create=True):
         else:
             self.output_channels = self.channels
 
-        self.dtype = get_output_dtype(self.model_config.output_dtype)
+        self.dtype = g.get_output_dtype(self.model_config.output_dtype)
 
         # g = Flow()
 
@@ -131,7 +118,6 @@ def __init__(self, yaml_config: str, create=True):
         )
         self.output_arrays = []
 
-
         output_shape = (
             np.array(self.idi_raw.shape)
             * np.array(self.input_voxel_size)
@@ -144,7 +130,7 @@ def __init__(self, yaml_config: str, create=True):
         for channel in self.output_channels:
             if create:
                 array = prepare_ds(
-                    DirectoryStore(self.output_path / channel/"s0"),
+                    DirectoryStore(self.output_path / channel / "s0"),
                     output_shape,
                     dtype=self.dtype,
                     chunk_shape=self.block_shape,
@@ -156,7 +142,7 @@ def __init__(self, yaml_config: str, create=True):
             else:
                 try:
                     array = open_ds(
-                        DirectoryStore(self.output_path / channel/"s0"),
+                        DirectoryStore(self.output_path / channel / "s0"),
                         "a",
                     )
                 except Exception as e:
@@ -247,7 +233,7 @@ def run(self):
 import subprocess
 
 
-def spawn_worker(name, yaml_config,charge_group,queue,ncpu=12):
+def spawn_worker(name, yaml_config, charge_group, queue, ncpu=12):
     def run_worker():
         subprocess.run(
             [
@@ -270,7 +256,7 @@ def run_worker():
                 "run",
                 "-y",
                 f"{yaml_config}",
-                "--client"
+                "--client",
             ]
         )
 

cellmap_flow/cli/cli.py

@@ -149,7 +149,6 @@ def script(script_path, data_path, queue, charge_group):
     help="The chargeback group to use when submitting",
     default=None,
 )
-
 def bioimage(model_path, data_path, edge_length_to_process, queue, charge_group):
     command = f"{SERVER_COMMAND} bioimage -m {model_path} -d {data_path} -e {edge_length_to_process}"
     base_name = model_path.split("/")[-1].split(".")[0]
@@ -188,7 +187,6 @@ def cellmap_model(config_folder, name, data_path, queue, charge_group):
     run(command, data_path, queue, charge_group, name)
 
 
-
 @cli.command()
 @click.option(
     "--script_path",

cellmap_flow/cli/multiple_cli.py

@@ -8,7 +8,7 @@
 import logging
 from cellmap_flow.utils.bsub_utils import start_hosts, SERVER_COMMAND
 from cellmap_flow.utils.neuroglancer_utils import generate_neuroglancer_url
-import cellmap_flow.globals as g
+from cellmap_flow.globals import g
 
 
 data_args = ["-d", "--data-path"]
@@ -158,11 +158,9 @@ def main():
                 else:
                     j += 1
             if not config_folder:
-                logger.error(
-                    "Missing -c for --celmmap-model sub-command."
-                )
+                logger.error("Missing -c for --cellmap-model sub-command.")
                 sys.exit(1)
-            models.append(CellMapModelConfig(config_folder, name=name,scale=scale))
+            models.append(CellMapModelConfig(config_folder, name=name, scale=scale))
             i = j
             continue
 

cellmap_flow/cli/multiple_yaml_cli.py

@@ -4,7 +4,8 @@
 from cellmap_flow.utils.bsub_utils import start_hosts, SERVER_COMMAND
 from cellmap_flow.utils.neuroglancer_utils import generate_neuroglancer_url
 from cellmap_flow.utils.config_utils import load_config, build_models
-import cellmap_flow.globals as g
+from cellmap_flow.globals import g
+
 
 logger = logging.getLogger(__name__)
 
@@ -24,10 +25,7 @@ def run_multiple(models, dataset_path, charge_group, queue):
 
         command = f"{SERVER_COMMAND} {model.command} -d {current_data_path}"
         start_hosts(
-            command,
-            job_name=model.name,
-            queue=queue,
-            charge_group=charge_group
+            command, job_name=model.name, queue=queue, charge_group=charge_group
         )
 
     generate_neuroglancer_url(dataset_path)
@@ -49,14 +47,14 @@ def main():
     config_path = sys.argv[1]
     config = load_config(config_path)
 
-    data_path = config['data_path']
-    charge_group = config['charge_group']
-    queue = config['queue']
+    data_path = config["data_path"]
+    charge_group = config["charge_group"]
+    queue = config["queue"]
 
     print("Data path:", data_path)
 
     # Build model configuration objects
-    models = build_models(config['models'])
+    models = build_models(config["models"])
 
     # For debugging, print each model config
     for model in models:

cellmap_flow/dashboard/app.py

@@ -22,7 +22,7 @@
     POSTPROCESS_DICT_KEY,
 )
 from cellmap_flow.models.run import update_run_models
-import cellmap_flow.globals as g
+from cellmap_flow.globals import g
 import numpy as np
 import time
 

cellmap_flow/globals.py

@@ -1,34 +1,128 @@
-jobs = []
+from cellmap_flow.norm.input_normalize import MinMaxNormalizer
+from cellmap_flow.post.postprocessors import DefaultPostprocessor
+from cellmap_flow.models.model_yaml import load_model_paths
+import os
+import threading
+import numpy as np
 
-models_config = []
 
-servers = []
+class Flow:
+    _instance = None
 
-raw = None
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(Flow, cls).__new__(cls)
+            cls._instance.jobs = []
+            cls._instance.models_config = []
+            cls._instance.servers = []
+            cls._instance.raw = None
+            cls._instance.input_norms = []  # or [MinMaxNormalizer(0, 255)]
+            cls._instance.postprocess = []
+            cls._instance.viewer = None
+            cls._instance.dataset_path = None
+            cls._instance.model_catalog = {}
+            # Uncomment and adjust if you want to load the model catalog:
+            # cls._instance.model_catalog = load_model_paths(
+            #     os.path.normpath(os.path.join(os.path.dirname(__file__), os.pardir, "models", "models.yaml"))
+            # )
+            cls._instance.queue = "gpu_h100"
+            cls._instance.charge_group = "cellmap"
+            cls._instance.neuroglancer_thread = None
+        return cls._instance
 
+    def to_dict(self):
+        return self.__dict__.items()
 
-from cellmap_flow.norm.input_normalize import MinMaxNormalizer
-from cellmap_flow.post.postprocessors import DefaultPostprocessor
+    def __repr__(self):
+        return f"Flow({self.__dict__})"
 
-# input_norms = [MinMaxNormalizer()]
-# postprocess = [DefaultPostprocessor(0,200,0,1)]
+    def __str__(self):
+        return f"Flow({self.__dict__})"
 
-input_norms = []
-postprocess = []
-viewer = None
+    def get_output_dtype(self):
+        dtype = np.float32
 
-dataset_path = None
+        if len(self.input_norms) > 0:
+            for norm in self.input_norms[::-1]:
+                if norm.dtype:
+                    dtype = norm.dtype
+                    break
 
+        if len(self.postprocess) > 0:
+            for postprocess in self.postprocess[::-1]:
+                if postprocess.dtype:
+                    dtype = postprocess.dtype
+                    break
 
-from cellmap_flow.models.model_yaml import load_model_paths
+        return dtype
 
-import os
-model_catalog = {}
-# model_catalog = load_model_paths(
-#   os.path.normpath(
-#      os.path.join(os.path.dirname(__file__), os.pardir, "models", "models.yaml")
-#  )
-# )
-
-queue = "gpu_h100"
-charge_group = "cellmap"
+    @classmethod
+    def run(
+        cls,
+        zarr_path,
+        model_configs,
+        queue="gpu_h100",
+        charge_group="cellmap",
+        input_normalizers=None,
+        post_processors=None,
+    ):
+
+        from cellmap_flow.utils.bsub_utils import start_hosts, SERVER_COMMAND
+        from cellmap_flow.utils.neuroglancer_utils import generate_neuroglancer_url
+
+        if input_normalizers is None:
+            input_normalizers = []
+        if post_processors is None:
+            post_processors = []
+
+        # Get the singleton instance (creates one if it doesn't exist)
+        instance = cls()
+        instance.queue = queue
+        instance.charge_group = charge_group
+        instance.dataset_path = zarr_path
+        instance.input_norms = input_normalizers
+        instance.postprocess = post_processors
+        instance.models_config = model_configs
+        instance.neuroglancer_thread = None
+
+        threads = []
+
+        for model_config in instance.models_config:
+            model_command = model_config.command
+            command = f"{SERVER_COMMAND} {model_command} -d {instance.dataset_path}"
+            print(f"Starting server with command: {command}")
+            thread = threading.Thread(
+                target=start_hosts,
+                args=(command, queue, charge_group, model_config.name),
+            )
+            thread.start()
+            threads.append(thread)
+
+        for thread in threads:
+            thread.join()
+
+        instance.neuroglancer_thread = threading.Thread(
+            target=generate_neuroglancer_url, args=(instance.dataset_path,)
+        )
+        instance.neuroglancer_thread.start()
+        # Optionally wait for the neuroglancer thread:
+        # instance.neuroglancer_thread.join()
+
+        print(f"*****Neuroglancer URL: {instance.dataset_path}")
+
+    @classmethod
+    def stop(cls):
+        instance = cls()
+        for job in instance.jobs:
+            print(f"Killing job {job.job_id}")
+            job.kill()
+        if instance.neuroglancer_thread is not None:
+            instance.neuroglancer_thread = None
+        instance.jobs = []
+
+    @classmethod
+    def delete(cls):
+        cls._instance = None
+
+
+g = Flow()

cellmap_flow/inferencer.py

@@ -3,22 +3,16 @@
 import torch
 from funlib.geometry import Coordinate
 import logging
-from cellmap_flow.utils.data import (
-    ModelConfig,
-    BioModelConfig,
-    DaCapoModelConfig,
-    ScriptModelConfig,
-    CellMapModelConfig,
-)
-import cellmap_flow.globals as g
-import neuroglancer
-from scipy import spatial
+from cellmap_flow.utils.data import ModelConfig
+
+from cellmap_flow.globals import g
 
 logger = logging.getLogger(__name__)
 
 
 def apply_postprocess(data, **kwargs):
     for pross in g.postprocess:
+        # logger.error(f"applying postprocess: {pross}")
         data = pross(data, **kwargs)
     return data
 
@@ -93,23 +87,13 @@ def optimize_model(self):
         self.model_config.config.model.eval()
 
     def process_chunk(self, idi, roi):
-        # if isinstance(self.model_config, BioModelConfig):
-        #     return self.process_chunk_bioimagezoo(idi, roi)
-        if (
-            isinstance(self.model_config, DaCapoModelConfig)
-            or isinstance(self.model_config, ScriptModelConfig)
-            or isinstance(self.model_config, BioModelConfig)
-          or isinstance(self.model_config, CellMapModelConfig)
+        # check if process_chunk is in self.config
+        if getattr(self.model_config.config, "process_chunk", None) and callable(
+            self.model_config.config.process_chunk
         ):
-            # check if process_chunk is in self.config
-            if getattr(self.model_config.config, "process_chunk", None) and callable(
-                self.model_config.config.process_chunk
-            ):
-                result = self.model_config.config.process_chunk(idi, roi)
-            else:
-                result = self.process_chunk_basic(idi, roi)
+            result = self.model_config.config.process_chunk(idi, roi)
         else:
-            raise ValueError(f"Invalid model config type {type(self.model_config)}")
+            result = self.process_chunk_basic(idi, roi)
 
         postprocessed = apply_postprocess(
             result,
@@ -131,4 +115,3 @@ def process_chunk_basic(self, idi, roi):
             use_half_prediction=self.use_half_prediction,
         )
         return result
-

cellmap_flow/models/run.py

@@ -1,4 +1,4 @@
-import cellmap_flow.globals as g
+from cellmap_flow.globals import g
 
 from cellmap_flow.utils.bsub_utils import start_hosts, SERVER_COMMAND
 from cellmap_flow.utils.web_utils import (