Visualise ellipses on annotated data (#297)

* Make notebook runnable with current data

* Make mask int16 rather than bool. Skip frames with no prompts. Rename vars

* Clarify

* Update notebook

* Update mask path

* Load frames too

Author: sfmig

notebooks/notebook_visualise_masks_from_zarr.py

@@ -7,11 +7,10 @@
 import matplotlib.pyplot as plt
 import napari
 import numpy as np
-import sparse
 import xarray as xr
-import zarr
 from ethology.io.annotations import load_bboxes
 from PIL import Image
+from skimage.measure import regionprops
 
 # %%
 # %matplotlib widget
@@ -27,26 +26,40 @@
 # %%%%%%%%%%%%%%%%%%%%%%%%%%
 # Helpers
 
+
 class ImageArrayLazy:
     """A lazy array for images in a list."""
 
     def __init__(self, img_paths):
         self.img_paths = sorted(img_paths)
-        # add image shape, assuming all have same as
+        # add image shape and dtype, assuming all have same as
         # first sample
-        sample = np.array(Image.open(img_paths[0]))  # H, W, C
+        sample = np.array(Image.open(self.img_paths[0]))  # H, W, C
         self.img_h, self.img_w, self.img_c = sample.shape
+        self.dtype = sample.dtype
 
     def __len__(self):
         return len(self.img_paths)
 
     def __getitem__(self, idx):
-        return np.array(Image.open(self.img_paths[idx]))
+        # napari indexes with a tuple, e.g. (frame, slice, slice)
+        frame_idx, *rest = idx if isinstance(idx, tuple) else (idx,)
+        if isinstance(frame_idx, slice):
+            frames = np.stack(
+                [np.array(Image.open(p)) for p in self.img_paths[frame_idx]]
+            )
+            return frames[(slice(None), *rest)]
+        return np.array(Image.open(self.img_paths[frame_idx]))[tuple(rest)]
 
     @property
     def shape(self):
+        """Shape of array (B, H, W, C)."""
         return (len(self.img_paths), self.img_h, self.img_w, self.img_c)
-        # B, H, W, C
+
+    @property
+    def ndim(self):
+        """Number of dimensions in array."""
+        return len(self.shape)
 
 
 def ellipses_from_labels(label_image):
@@ -63,6 +76,7 @@ def ellipses_from_labels(label_image):
         Endpoint pairs for the major axis (y, x).
     minor_axes : list of (2, 2) arrays
         Endpoint pairs for the minor axis (y, x).
+
     """
     ellipse_corners = []
     major_axes = []
@@ -124,12 +138,12 @@ def ellipses_from_labels(label_image):
 
 # %%%%%%%%%%%%%%%%%%%%%%%
 # Load masks
-zarr_root = zarr.open(
-    DATA_DIR / f"{LABEL_NAME} masks.zarr",
-    mode="r",
-)
+# ID-encoded masks (int16): 0 = background, k + 1 = the k-th box
+# in the frame. Set the store name to the output of
+# scripts/generate_masks_from_bboxes.py
+MASKS_ZARR = DATA_DIR / "annotations" / "masks_20261009_133849.zarr"
 
-mask_da_array = da.from_zarr(zarr_root["masks"])
+mask_da_array = da.from_zarr(MASKS_ZARR, mode="r")  # (image_id, H, W)
 
 
 # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
@@ -138,14 +152,20 @@ def ellipses_from_labels(label_image):
 # CONS: can be slow to compute
 n_ids = ds_bboxes.sizes["id"]
 
+
 # for one task per chunk:
 def _apply_one_hot_encoding_for_IDs(block, n_ids):
-    """Transform a block of shape (chunk_frames, H, W) 
+    """Apply one hot encoding to ID-encoded block.
+
+    Transform an ID-encoded block of shape (chunk_frames, H, W)
     into a boolean array of shape (chunk_frames, n_ids, H, W).
 
+    Label k + 1 in the block is the mask of the box at position k
+    along the `id` dimension.
+
     This is most efficient if chunk_frames=1.
     """
-    labels = np.arange(1, n_ids + 1)[None, :, None, None] # (1, n_ids, 1, 1)
+    labels = np.arange(1, n_ids + 1)[None, :, None, None]  # (1, n_ids, 1, 1)
     # compare every pixel against every label simultaneously
     # via broadcasting (i.e. for each pixel, we have a 1-hot encoding
     # vector indicate which label/id it is)
@@ -154,16 +174,16 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids):
 
 # Apply expand labels chunk by chunk
 # we define one chunk per image
-mask_da_array = mask_da_array.rechunk({0: 1}) 
+mask_da_array = mask_da_array.rechunk({0: 1})
 mask_4d = mask_da_array.map_blocks(
     _apply_one_hot_encoding_for_IDs,
     n_ids=n_ids,
-    new_axis=1, # position of new_axis id
-    chunks=( # sizes of output chunks (if diff from input)
-        mask_da_array.chunks[0], # same as input chunk axis=0
-        (n_ids,),   # ids
-        mask_da_array.chunks[1], # same as input chunk axis=1
-        mask_da_array.chunks[2], # same as input chunk axis=2
+    new_axis=1,  # position of new_axis id
+    chunks=(  # sizes of output chunks (if diff from input)
+        mask_da_array.chunks[0],  # same as input chunk axis=0
+        (n_ids,),  # ids
+        mask_da_array.chunks[1],  # same as input chunk axis=1
+        mask_da_array.chunks[2],  # same as input chunk axis=2
     ),
     dtype=bool,
 )  # (image_id, id, img_h, img_w)
@@ -201,32 +221,32 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids):
 )
 ax.contour(single_mask, levels=[0.5], colors="red", linewidths=0.5)
 
-# plot all masks in one frame (boolean masks, all same color!)
-all_masks = ds_bboxes.masks_bool.sel(image_id=image_id).any(dim="id")
+# plot all masks in one frame (one color per ID)
+label_frame = mask_da_array[image_id].compute()
 ax.imshow(
-    all_masks,
+    label_frame,
     cmap="turbo",
-    alpha=all_masks.astype(float) * 0.5,
+    alpha=(label_frame > 0) * 0.5,
 )
 
 # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 # Visualise frames and masks in napari viewer
 mask_array = mask_da_array
 viewer = napari.Viewer()
 
-# viewer.add_image(np.asarray(image_array).moveaxis(0, -1, 1, 2), name="image")
-
-viewer.add_labels(np.asarray(mask_array), name=f"{LABEL_NAME} masks")
+# both layers are lazy, so napari only loads the frame displayed
+viewer.add_image(image_array, name="frames", rgb=True)
+viewer.add_labels(mask_array, name=f"{LABEL_NAME} masks")
 
 # %%
 # build ellipses and axis lines for napari Shapes layers
 all_ellipses = []
 all_major_axes = []
 all_minor_axes = []
 for frame_idx in range(mask_array.shape[0]):
-    ellipse_corners, major_axes, minor_axes = ellipses_from_labels(
-        np.asarray(mask_array[frame_idx])
-    )
+    # masks are already ID-encoded, so each ID is one region
+    label_frame = np.asarray(mask_array[frame_idx])
+    ellipse_corners, major_axes, minor_axes = ellipses_from_labels(label_frame)
     for corners in ellipse_corners:
         # Prepend frame index as first column for nD shapes
         nd = np.column_stack(
@@ -266,3 +286,5 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids):
     edge_width=4,
     name="major axes",
 )
+
+# %%

scripts/generate_masks_from_bboxes.py

@@ -3,6 +3,9 @@
 Pass bounding boxes as prompts to SAM2, predict masks in those regions,
 and save as an ID-encoded zarr store.
 
+This script is a one-off tool to run on annotated frames whose mask labels
+are only meaningful within a frame.
+
 Expected data directory structure::
 
     <data_dir>/
@@ -117,8 +120,8 @@ def create_mask_zarr(
         path_to_zarr,
         mode="w",
         shape=(n_images, image_h, image_w),
-        dtype="bool",
-        fill_value=False,
+        dtype="int16",
+        fill_value=0,
         chunks=(1, image_h, image_w),
     )
 
@@ -140,59 +143,86 @@ def predict_masks_across_images(
         1:3
     ]  # ds_bboxes.attrs["image_array"].shape[1:3]
 
-    for idx in range(0, n_frames, img_batch_size):
-        actual_batch_size = min(img_batch_size, n_frames - idx)
+    for batch_start in range(0, n_frames, img_batch_size):
+        actual_batch_size = min(img_batch_size, n_frames - batch_start)
 
         # Integer array (B, H, W):
         # 0 = background, 1+ = object instance IDs
-        id_mask_batch = np.zeros(
+        label_mask_batch = np.zeros(
             (actual_batch_size, img_h, img_w), dtype=np.int16
         )
 
-        # Compute embeddings for image batch
-        image_batch = [image_array[idx + i] for i in range(actual_batch_size)]
-        model.set_image_batch(image_batch)
-
         # List of (N_i, 4) arrays, one per image
         boxes_batch = [
             bbox_prompts_per_frame_idx[f_i]
-            for f_i in range(idx, idx + actual_batch_size)
+            for f_i in range(batch_start, batch_start + actual_batch_size)
         ]
 
-        # Predict batch of masks — list of (N, 1, H, W) arrays
-        masks_batch, _scores_batch, _ = model.predict_batch(
-            box_batch=boxes_batch,
-            multimask_output=False,
-        )
+        # Only pass frames with at least one box to SAM2;
+        # frames without boxes are left as background (0)
+        # idcs_rel = indices within the batch
+        idcs_rel_with_boxes = []
+        for i, boxes in enumerate(boxes_batch):
+            if len(boxes):
+                idcs_rel_with_boxes.append(i)
+            else:
+                print(f"Frame {batch_start + i}: no boxes, skipped")
+
+        if idcs_rel_with_boxes:
+            # Compute embeddings for images with boxes
+            image_batch = [
+                image_array[batch_start + i] for i in idcs_rel_with_boxes
+            ]
+            model.set_image_batch(image_batch)
 
-        # Convert boolean masks to ID-encoded masks
-        # (higher ID wins in overlap)
-        for idx_rel_batch in range(actual_batch_size):
-            masks_one_frame = masks_batch[idx_rel_batch].squeeze(
-                axis=1
-            )  # (N, H, W)
+            # List of (N_i, 4) arrays, one per image with boxes
+            box_batch = [boxes_batch[i] for i in idcs_rel_with_boxes]
 
-            n_objects = masks_one_frame.shape[0]
-            obj_ids = np.arange(1, n_objects + 1, dtype=np.int16)[
-                :, None, None
-            ]
-            id_mask_batch[idx_rel_batch] = (masks_one_frame * obj_ids).max(
-                axis=0
+            # Predict batch of masks — list of (N, 1, H, W) arrays,
+            # it has one per mask image for each frame **with boxes**
+            masks_batch, _scores_batch, _ = model.predict_batch(
+                box_batch=box_batch,
+                multimask_output=False,
             )
 
-            print(
-                f"Frame {idx + idx_rel_batch}: "
-                f"{n_objects} masks / "
-                f"{boxes_batch[idx_rel_batch].shape[0]}"
-                " boxes"
-            )
+            # Convert boolean masks to ID-encoded masks
+            # (higher ID wins in overlap)
+            # idcs_rel = indices within the batch
+            for idx_rel_batch, masks in zip(
+                idcs_rel_with_boxes,
+                masks_batch,
+                strict=True,
+            ):
+                # SAM2 squeezes the box axis when N == 1, returning
+                # (1, H, W) instead of (N, 1, H, W); we use reshape
+                # to handle both
+                masks_one_frame = masks.reshape(-1, img_h, img_w)  # (N, H, W)
+                n_objects = masks_one_frame.shape[0]
+                obj_ids = np.arange(1, n_objects + 1, dtype=np.int16)[
+                    :, None, None
+                ]
+
+                # Add results to batch array (B,H,W)
+                label_mask_batch[idx_rel_batch] = (
+                    masks_one_frame * obj_ids
+                ).max(axis=0)
+
+                # log
+                print(
+                    f"Frame {batch_start + idx_rel_batch}: "
+                    f"{n_objects} masks / "
+                    f"{boxes_batch[idx_rel_batch].shape[0]}"
+                    " boxes"
+                )
 
         # Save id_mask to zarr store
-        output_zarr[idx : idx + actual_batch_size] = id_mask_batch
+        output_zarr[batch_start : batch_start + actual_batch_size] = (
+            label_mask_batch
+        )
 
         # Mark frames as annotated in zarr store attributes
         annotated = set(output_zarr.attrs.get("annotated_frames", []))
-        annotated.update(range(idx, idx + actual_batch_size))
+        annotated.update(range(batch_start, batch_start + actual_batch_size))
         output_zarr.attrs["annotated_frames"] = sorted(annotated)