Model fitting works lazily

src/cellcanvas/_app/main_app.py

@@ -30,6 +30,11 @@ def __init__(
             extra_logging=self.verbose,
         )
 
+
+    def update_data_manager(self, data: DataManager):
+        self.data = data
+        self.semantic_segmentor.update_data_manager(data)
+
     @property
     def mode(self) -> AppMode:
         return self._mode

src/cellcanvas/data/data_manager.py

@@ -3,6 +3,7 @@
 import numpy as np
 from napari.utils.events.containers import SelectableEventedList
 from zarr import Array
+import dask.array as da
 
 from cellcanvas.data.data_set import DataSet
 
@@ -15,6 +16,7 @@ def __init__(self, datasets: Optional[List[DataSet]] = None):
             datasets = [datasets]
         self.datasets = SelectableEventedList(datasets)
 
+    # Normal version
     def get_training_data(self) -> Tuple[Array, Array]:
         """Get the pixel-wise semantic segmentation training data for datasets.
 
@@ -30,23 +32,26 @@ def get_training_data(self) -> Tuple[Array, Array]:
         features = []
         labels = []
         for dataset in self.datasets:
-            # get the features and labels
-            # todo make lazier
-            dataset_features = np.asarray(dataset.concatenated_features)
-            dataset_labels = np.asarray(dataset.labels)
-
-            # reshape the data
-            dataset_labels = dataset_labels.flatten()
-            reshaped_features = dataset_features.reshape(
-                -1, dataset_features.shape[-1]
-            )
-
-            # Filter features where labels are greater than 0
-            valid_labels = dataset_labels > 0
-            filtered_features = reshaped_features[valid_labels, :]
-            filtered_labels = dataset_labels[valid_labels] - 1  # Adjust labels
+            dataset_features = da.asarray(dataset.concatenated_features)
+            dataset_labels = da.asarray(dataset.labels)
+            # Flatten labels for boolean indexing
+            flattened_labels = dataset_labels.flatten()
+
+            # Compute valid_indices based on labels > 0
+            valid_indices = da.nonzero(flattened_labels > 0)[0].compute()
+
+            # Flatten only the spatial dimensions of the dataset_features while preserving the feature dimension
+            c, h, w, d = dataset_features.shape
+            reshaped_features = dataset_features.reshape(c, h * w * d)
+
+            # We need to apply valid_indices for each feature dimension separately
+            filtered_features_list = [da.take(reshaped_features[i, :], valid_indices, axis=0) for i in range(c)]
+            filtered_features = da.stack(filtered_features_list, axis=1)
+
+            # Adjust labels
+            filtered_labels = flattened_labels[valid_indices] - 1
 
             features.append(filtered_features)
             labels.append(filtered_labels)
-
-        return np.concatenate(features), np.concatenate(labels)
+            
+        return da.concatenate(features), da.concatenate(labels)

src/cellcanvas/data/data_set.py

@@ -6,6 +6,8 @@
 import zarr
 from zarr import Array
 
+from ome_zarr.io import ZarrLocation
+from ome_zarr.reader import Multiscales
 
 @dataclass
 class DataSet:
@@ -62,7 +64,11 @@ def from_paths(
                 dimension_separator=".",
             )
         else:
-            labels = zarr.open(labels_path, "a")
+            if Multiscales.matches(ZarrLocation(labels_path)):
+                labels = zarr.open(os.path.join(labels_path, "0"),
+                                   "a")
+            else:
+                labels = zarr.open(labels_path, "a")
 
         # get the segmentation
         if (not os.path.isdir(segmentation_path)) and make_missing_datasets:
@@ -83,3 +89,57 @@ def from_paths(
             labels=labels,
             segmentation=segmentation,
         )
+
+    @classmethod
+    def from_stores(
+        cls,
+        image_store,
+        features_store,
+        labels_store,
+        segmentation_store,
+    ):
+        """Create a DataSet from a set of paths.
+
+        todo: add ability to create missing labels/segmentations
+        """
+
+        # TODO rewrite this to copy everything to be local
+
+        # get the image
+        # TODO fix hardcoded scale for pickathon
+        image = zarr.open(zarr.storage.LRUStoreCache(image_store, None), "r")["0"]
+
+        # get the features
+        features = {"features": zarr.open(zarr.storage.LRUStoreCache(features_store, None), "r")}
+
+        group_name = "labels"
+
+        # get the labels
+        labels = zarr.open_group(zarr.storage.LRUStoreCache(labels_store, None),
+                                 mode="a")
+        if group_name in labels:
+            labels = labels[group_name]
+        else:
+            labels = labels.create_dataset(group_name,
+                                           shape=image.shape,
+                                           dtype="i4")
+
+        # get the segmentation
+        segmentation = zarr.open_group(zarr.storage.LRUStoreCache(segmentation_store, None),
+                                       mode="a")
+        if group_name in segmentation:
+            segmentation = segmentation[group_name]
+        else:
+            segmentation = segmentation.create_dataset(group_name,
+                                                       shape=image.shape,
+                                                       dtype="i4")
+
+        # TODO start a background thread that triggers downloads of the zarrs
+
+        return cls(
+            image=image,
+            features=features,
+            labels=labels,
+            segmentation=segmentation,
+        )
+