arbitrary number of imaging layers in zarr

.gitignore

@@ -166,6 +166,8 @@ cython_debug/
 
 # vscode files
 .history
+.vscode/
+CLAUDE.md
 
 # Ignore examples files
 examples/*/*

pyproject.toml

@@ -39,10 +39,10 @@ extra-dependencies = [
   "mypy>=1.0.0",
 ]
 [tool.hatch.envs.types.scripts]
-check = "mypy --install-types --non-interactive {args:trackedit tests}"
+check = "mypy --install-types --non-interactive {args:trackedit trackedit/_tests}"
 
 [tool.coverage.run]
-source_pkgs = ["trackedit", "tests"]
+source_pkgs = ["trackedit", "trackedit/_tests"]
 branch = true
 parallel = true
 omit = [
@@ -51,7 +51,7 @@ omit = [
 
 [tool.coverage.paths]
 trackedit = ["trackedit", "*/trackedit/trackedit"]
-tests = ["tests", "*/trackedit/tests"]
+tests = ["trackedit/_tests", "*/trackedit/trackedit/_tests"]
 
 [tool.coverage.report]
 exclude_lines = [
@@ -90,6 +90,7 @@ pyqt = ">=5.15.9,<6"
 numpy = "<2.2"
 pre-commit = ">=4.1.0,<5"
 dask = ">=2025.2.0,<2026"
+pytest = ">=8.4.2,<9"
 
 [tool.pixi.feature.test.dependencies]
 pytest = "*"

scripts/script_cropped_fov.py

@@ -27,7 +27,7 @@
 
 # OPTIONAL: imaging data
 imaging_zarr_file = (
-    "/hpc/projects/intracellular_dashboard/organelle_dynamics/rerun/"
+    "/hpc/projects/intracellular_dashboard/organelle_dynamics/"
     "2025_07_24_A549_SEC61_TOMM20_G3BP1_ZIKV/2-assemble/"
     "2025_07_24_A549_SEC61_TOMM20_G3BP1_ZIKV.zarr/A/1/000000/"
 )

trackedit/DatabaseHandler.py

@@ -58,6 +58,7 @@ def __init__(
         image_translate: tuple = None,
         coordinate_filters: list = None,
         default_start_annotation: int = None,  # Make it optional
+        imaging_layer_names: list = None,
     ):
 
         # inputs
@@ -75,8 +76,11 @@ def __init__(
         self.imaging_channel = imaging_channel
         self.image_z_slice = image_z_slice
         self.image_translate = image_translate
-        self.imaging_flag = True if self.imaging_zarr_file is not None else False
+        self.imaging_flag = (
+            True if self.imaging_zarr_file and self.imaging_zarr_file != "" else False
+        )
         self.coordinate_filters = coordinate_filters
+        self.imaging_layer_names = imaging_layer_names
 
         # Filenames / directories
         self.extension_string = ""
@@ -175,6 +179,7 @@ def __init__(
                 channel=self.imaging_channel,
                 time_window=self.time_window,
                 image_z_slice=self.image_z_slice,
+                imaging_layer_names=self.imaging_layer_names,
             )
         self.df_full = self.db_to_df(entire_database=True)
 

trackedit/TrackEditClass.py

@@ -72,27 +72,28 @@ def __init__(
             color_dict=self.databasehandler.color_mapping
         )
 
-        # Add Imaging layer to viewer
+        # Add Imaging layers to viewer
         if self.databasehandler.imaging_flag:
-            layer_nuc = self.viewer.add_image(
-                self.databasehandler.imagingArray.nuclear,
-                name="im_nuclear",
-                colormap="green",
-                scale=self.databasehandler.scale,
-                visible=False,
-                translate=self.databasehandler.image_translate,
-            )
-            layer_nuc.reset_contrast_limits()
-            layer_mem = self.viewer.add_image(
-                self.databasehandler.imagingArray.membrane,
-                name="im_membrane",
-                colormap="red",
-                opacity=0.5,
-                scale=self.databasehandler.scale,
-                visible=False,
-                translate=self.databasehandler.image_translate,
-            )
-            layer_mem.reset_contrast_limits()
+            # Define colormap sequence
+            colormaps = ["green", "red", "blue", "magenta", "cyan", "yellow"]
+
+            for i, layer_name in enumerate(
+                self.databasehandler.imagingArray.layer_names
+            ):
+                channel_data = self.databasehandler.imagingArray.get_channel_data(i)
+                colormap = colormaps[i % len(colormaps)]
+                opacity = 1.0 if i == 0 else 0.5  # First layer full opacity, others 0.5
+
+                layer = self.viewer.add_image(
+                    channel_data,
+                    name=f"im_{layer_name}",
+                    colormap=colormap,
+                    opacity=opacity,
+                    scale=self.databasehandler.scale,
+                    visible=False,
+                    translate=self.databasehandler.image_translate,
+                )
+                layer.reset_contrast_limits()
 
         tabwidget_bottom = QTabWidget()
         tabwidget_bottom.addTab(self.TreeWidget, "TreeWidget")
@@ -188,12 +189,11 @@ def update_chunk_from_button(self, direction: str):
         self.databasehandler.set_time_chunk(new_chunk)
         self.update_hierarchy_layer()
         if self.databasehandler.imaging_flag:
-            self.viewer.layers[
-                "im_nuclear"
-            ].data = self.databasehandler.imagingArray.nuclear
-            self.viewer.layers[
-                "im_membrane"
-            ].data = self.databasehandler.imagingArray.membrane
+            for i, layer_name in enumerate(
+                self.databasehandler.imagingArray.layer_names
+            ):
+                channel_data = self.databasehandler.imagingArray.get_channel_data(i)
+                self.viewer.layers[f"im_{layer_name}"].data = channel_data
 
         self.add_tracks()
 
@@ -244,12 +244,11 @@ def update_chunk_from_frame(self, frame: int):
         self.databasehandler.set_time_chunk(new_chunk)
         self.update_hierarchy_layer()
         if self.databasehandler.imaging_flag:
-            self.viewer.layers[
-                "im_nuclear"
-            ].data = self.databasehandler.imagingArray.nuclear
-            self.viewer.layers[
-                "im_membrane"
-            ].data = self.databasehandler.imagingArray.membrane
+            for i, layer_name in enumerate(
+                self.databasehandler.imagingArray.layer_names
+            ):
+                channel_data = self.databasehandler.imagingArray.get_channel_data(i)
+                self.viewer.layers[f"im_{layer_name}"].data = channel_data
 
         # Update tracks if chunks are different OR if this was triggered by a Tmax change
         # TODO: not sure this is the best way to do this

trackedit/arrays/ImagingArray.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import List, Optional, Tuple
 
 import dask.array as da
 
@@ -12,13 +12,17 @@ def __init__(
         channel: str = "0/4/0/0",
         time_window: Tuple[int, int] = (0, 104),
         image_z_slice: int = None,
+        imaging_layer_names: Optional[List[str]] = None,
     ):
         """
         Initialize the image array from a zarr file.
 
         Args:
             imaging_zarr_file: Path to the zarr file
             channel: Channel path within the zarr file (default: '0/4/0/0')
+            time_window: Time window for the image stack
+            image_z_slice: Optional z-slice to extract
+            imaging_layer_names: Names for imaging channels. If None, defaults to ['nuclear', 'membrane'] for 2 channels
         """
         # Load the full stack using dask
         self._full_stack = da.from_zarr(imaging_zarr_file, component=channel)
@@ -32,24 +36,56 @@ def __init__(
         # Set initial stack based on full time window
         self._update_stack()
 
+        # Detect number of channels and set layer names
+        self.n_channels = self._stack.shape[1] if len(self._stack.shape) > 1 else 1
+
+        # Set layer names with backward compatibility
+        if imaging_layer_names is None:
+            if self.n_channels == 2:
+                self.layer_names = ["nuclear", "membrane"]
+            else:
+                self.layer_names = [f"channel_{i}" for i in range(self.n_channels)]
+        else:
+            if len(imaging_layer_names) != self.n_channels:
+                raise ValueError(
+                    f"Number of layer names ({len(imaging_layer_names)}) "
+                    f"doesn't match number of zarr channels ({self.n_channels})"
+                )
+            self.layer_names = imaging_layer_names.copy()
+
     def _update_stack(self):
         """Update the stack based on current time window"""
         self._stack = self._full_stack[self._time_window[0] : self._time_window[1]]
 
+    def get_channel_data(self, channel_idx: int) -> da.Array:
+        """Get data for a specific channel by index"""
+        if channel_idx >= self.n_channels:
+            raise ValueError(
+                f"Channel index {channel_idx} >= number of channels ({self.n_channels})"
+            )
+        return self._stack[:, channel_idx]
+
+    def get_channel_by_name(self, name: str) -> da.Array:
+        """Get channel data by name"""
+        if name not in self.layer_names:
+            raise ValueError(f"Channel name '{name}' not found in {self.layer_names}")
+        idx = self.layer_names.index(name)
+        return self.get_channel_data(idx)
+
     @property
     def nuclear(self) -> da.Array:
-        """Get the nuclear channel data"""
-        return self._stack[:, 0]
+        """Get the nuclear channel data (backward compatibility)"""
+        return self.get_channel_by_name("nuclear")
 
     @property
     def membrane(self) -> da.Array:
-        """Get the membrane channel data"""
-        return self._stack[:, 1]
+        """Get the membrane channel data (backward compatibility)"""
+        return self.get_channel_by_name("membrane")
 
     @property
     def shape(self) -> Tuple[int, ...]:
-        """Get the shape of a single channel (identical for nuclear and membrane)"""
-        return self._stack[:, 0].shape
+        """Get the shape of a single channel"""
+        return self.get_channel_data(0).shape
 
     @property
     def time_window(self) -> Tuple[int, int]:

trackedit/run.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import napari
 
@@ -42,6 +42,7 @@ def run_trackedit(
     annotation_mapping: Optional[dict] = None,
     coordinate_filters: Optional[list] = None,
     default_start_annotation: Optional[int] = None,
+    imaging_layer_names: Optional[List[str]] = None,
 ) -> Tuple[napari.Viewer, TrackEditClass]:
     """Run TrackEdit on a database file.
 
@@ -61,6 +62,7 @@ def run_trackedit(
         viewer: Optional existing napari viewer
         flag_show_hierarchy: Show hierarchy in the viewer
         annotation_mapping: Mapping of annotation ids to names and colors
+        imaging_layer_names: Names for imaging layers. If None, defaults to ['nuclear', 'membrane'] for 2 channels
 
     Returns:
         Tuple[napari.Viewer, TrackEditClass]: The viewer instance and TrackEdit instance
@@ -87,6 +89,7 @@ def run_trackedit(
         image_translate=image_translate,
         coordinate_filters=coordinate_filters,
         default_start_annotation=default_start_annotation,
+        imaging_layer_names=imaging_layer_names,
     )
 
     # overwrite some motile functions