diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml index 6823d09c..78173140 100644 --- a/.github/workflows/test_and_deploy.yml +++ b/.github/workflows/test_and_deploy.yml @@ -49,7 +49,7 @@ jobs: shell: bash -l {0} run: | sudo apt-get update - sudo apt-get install libglapi-mesa libegl-mesa0 libegl1 libopengl0 libgl1 + sudo apt-get install libegl-mesa0 libegl1 libopengl0 libgl1 - name: Run all tests uses: neuroinformatics-unit/actions/test@v3 with: @@ -76,7 +76,7 @@ jobs: - name: Install OpenGL libraries run: | sudo apt-get update - sudo apt-get install libglapi-mesa libegl-mesa0 libegl1 libopengl0 libgl1 + sudo apt-get install libegl-mesa0 libegl1 libopengl0 libgl1 # as in bash_scripts/run_mask_array.sh: install torch (and setuptools) # first, then build SAM2 against them without build isolation - name: Install crabs, then SAM2 diff --git a/notebooks/notebook_visualise_masks_from_zarr.py b/notebooks/notebook_visualise_masks_from_zarr.py index ca96b416..691df778 100644 --- a/notebooks/notebook_visualise_masks_from_zarr.py +++ b/notebooks/notebook_visualise_masks_from_zarr.py @@ -3,212 +3,218 @@ # %% from pathlib import Path +import dask import dask.array as da import matplotlib.pyplot as plt import napari import numpy as np import xarray as xr -from ethology.io.annotations import load_bboxes +from matplotlib.patches import Ellipse from PIL import Image -from skimage.measure import regionprops # %% # %matplotlib widget # %%%%%%%%%%%%%%%%%%%%%%%%%%%%% -# September groundtruth data -DATA_DIR = Path("/Users/sofia/swc/CrabLabels/sep2023-full") -IMAGES_DIR = DATA_DIR / "frames" -ANNOTATIONS_FILE = DATA_DIR / "annotations" / "VIA_JSON_combined_coco_gen.json" +# Groundtruth data -LABEL_NAME = "crab" - -# %%%%%%%%%%%%%%%%%%%%%%%%%% -# Helpers +# map dataset to the output store of scripts/generate_masks_from_bboxes.py +# (stores named masks_.zarr, written before the bboxes and mask +# parameters were added, cannot be read here; re-run the script) +map_month_to_zarr = { + "sep": "bboxes_and_masks_20261026_123937.zarr", + "aug": "bboxes_and_masks_20261026_125630.zarr", +} -class ImageArrayLazy: - """A lazy array for images in a list.""" - - def __init__(self, img_paths): - self.img_paths = sorted(img_paths) - # add image shape and dtype, assuming all have same as - # first sample - 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 +# Select input data (sept/aug) +MONTH = "aug" +DATA_DIR = Path.home() / Path( + f"swc/project_crabs/data/CrabLabels/{MONTH}2023-full" +) +ZARR_STORE = DATA_DIR / "annotations" / map_month_to_zarr[MONTH] - def __len__(self): - return len(self.img_paths) - def __getitem__(self, 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)] +# %%%%%%%%%%%%%%%%%%%%%%%%%% +# Helpers - @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) - @property - def ndim(self): - """Number of dimensions in array.""" - return len(self.shape) +def _read_image(path): + """Read an image file as a numpy array of shape (H, W, C).""" + return np.array(Image.open(path)) -def ellipses_from_labels(label_image): - """Return ellipse corners, major-axis lines, and minor-axis lines. +def ellipse_shapes_from_ds(ds): + """Return napari ellipses, major axes and minor axes for every mask. - Each ellipse is represented by the 4 corners of its oriented bounding - rectangle (for napari's built-in ``"ellipse"`` shape type). + In napari, the ellipse shape is represented by the 4 corners of its + oriented bounding rectangle. The major and minor axes are represented + by their 2 endpoints, which have coordinates (image index, y, x). Returns ------- - ellipse_corners : list of (4, 2) arrays - Corners in (y, x) order. - major_axes : list of (2, 2) arrays - Endpoint pairs for the major axis (y, x). - minor_axes : list of (2, 2) arrays - Endpoint pairs for the minor axis (y, x). + For N masks: + - ellipse_corners : (N, 4, 3) array + - major_axes : (N, 2, 3) array + - minor_axes : (N, 2, 3) array - """ - ellipse_corners = [] - major_axes = [] - minor_axes = [] - for prop in regionprops(label_image): - cy, cx = prop.centroid - a = prop.major_axis_length / 2 - b = prop.minor_axis_length / 2 - theta = prop.orientation - # scikit-image orientation: angle of the major axis - # measured counter-clockwise from the row (y) axis - # direction vectors in (y, x) - dir_major = np.array([np.cos(theta), np.sin(theta)]) - dir_minor = np.array([-np.sin(theta), np.cos(theta)]) - - center = np.array([cy, cx]) - - # 4 corners of the oriented bounding rectangle - corners = np.array( - [ - center + a * dir_major + b * dir_minor, - center + a * dir_major - b * dir_minor, - center - a * dir_major - b * dir_minor, - center - a * dir_major + b * dir_minor, - ] - ) - ellipse_corners.append(corners) + The trailing 3 dimension corresponds to the coordinates + (image index, y, x). - # axis endpoints - major_axes.append( - np.array([center - a * dir_major, center + a * dir_major]) - ) - minor_axes.append( - np.array([center - b * dir_minor, center + b * dir_minor]) - ) - - return ellipse_corners, major_axes, minor_axes + """ + # Get ellipses data from dataset + + # (image_id, id) positions of masks + image_idcs, mask_idcs = np.nonzero(ds.mask_area.notnull().values) + + # mask centroid + cx = ds.mask_centroid.sel(space="x").values[image_idcs, mask_idcs] + cy = ds.mask_centroid.sel(space="y").values[image_idcs, mask_idcs] + c = np.stack([cy, cx], axis=-1) + + # semi-axes lengths + major_semi_length = ( + ds.ellipse_major_axis_length.values[image_idcs, mask_idcs][:, None] / 2 + ) + minor_semi_length = ( + ds.ellipse_minor_axis_length.values[image_idcs, mask_idcs][:, None] / 2 + ) + + # ellipse orientation wrt x axis image coords + theta = ds.ellipse_major_axis_theta.values[image_idcs, mask_idcs] + + # axes unit vectors in (y, x) space + major_unit_vector = np.stack([np.sin(theta), np.cos(theta)], axis=-1) + minor_unit_vector = np.stack([np.cos(theta), -np.sin(theta)], axis=-1) + + # semi-axis vectors in (y, x) + major_semi_vector = major_semi_length * major_unit_vector + minor_semi_vector = minor_semi_length * minor_unit_vector + + # corners of the ellipse's oriented bounding box + corners = np.stack( + [ + c + major_semi_vector + minor_semi_vector, + c + major_semi_vector - minor_semi_vector, + c - major_semi_vector - minor_semi_vector, + c - major_semi_vector + minor_semi_vector, + ], + axis=1, + ) # (N, 4, 2) + + # axes segments + major_axis_endpoints = np.stack( + [c - major_semi_vector, c + major_semi_vector], axis=1 + ) # (N, 2, 2) + minor_axis_endpoints = np.stack( + [c - minor_semi_vector, c + minor_semi_vector], axis=1 + ) # (N, 2, 2) + + def add_image_idx(points): + """Add image index coordinate to the input `points` array. + + The input `points` array is of shape (N, n_points, 2), + where N is the number of masks, n_points is the number + of points required to define the shape (n_points=4 for + corners, n_points=2 for axes segments), and 2 are the + spatial dimensions (y,x). + """ + # Broadcast image_idcs (N,) to (N, n_points, 1) to match + # the first two dimensions of the points array. + image_idcs_b = np.broadcast_to( + image_idcs[:, None, None], # (N,) --> (N, 1, 1) + (*points.shape[:2], 1), # (N, n_points, 1) + ) # (N, n_points, 1) + + # Concatenate image_idcs along the last axis (coordinates) + # (N, n_points, 1) + (N, n_points, 2) --> (N, n_points, 3) + return np.concatenate([image_idcs_b, points], axis=-1) + + return ( + add_image_idx(corners), + add_image_idx(major_axis_endpoints), + add_image_idx(minor_axis_endpoints), + ) # %%%%%%%%%%%%%%%%%%%%%% -# Read groundtruth as ethology annotation dataset -ds_bboxes = load_bboxes.from_files( - ANNOTATIONS_FILE, - format="COCO", - images_dirs=IMAGES_DIR, -) - -# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -# Load ground truth images as lazy array - -# build lazy image array -png_files = sorted(ds_bboxes.attrs["images_directories"].glob("*.png")) -image_array = ImageArrayLazy(png_files) -print(image_array.shape) - -# add as attribute -ds_bboxes.attrs["image_array"] = image_array - - -# %%%%%%%%%%%%%%%%%%%%%%% -# Load masks -# 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" +# Read bboxes, masks and geometric parameters from zarr store +# - labels: ID-encoded masks (uint16), 0 = background +# - label_of: label k + 1 is the mask for the box at position k along `id` +# - mask_* and ellipse_*: parameters per mask, on (image_id, id) +ds = xr.open_zarr(ZARR_STORE) -mask_da_array = da.from_zarr(MASKS_ZARR, mode="r") # (image_id, H, W) - - -# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -# Load as 4D boolean dask array with id dim -# Each slice [image_id, id] is a boolean mask for that object in that frame -# CONS: can be slow to compute -n_ids = ds_bboxes.sizes["id"] +print(ds.data_vars) +# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% +# Load ground truth images into dataset +# as a lazy dask array +# (image files are aligned with the image_id coordinate) + +# zarr stores attrs as JSON, so the image_id keys read back as str, +# which we need to cast to int here +map_id_to_filename = { + int(k): v for k, v in ds.attrs["map_image_id_to_filename"].items() +} +images_dir = Path(ds.attrs["images_directories"]) +img_paths = [ + images_dir / map_id_to_filename[img_id] for img_id in ds.image_id.values +] + +# assume all images have the same shape and dtype as the first one +sample = _read_image(img_paths[0]) # H, W, C +# define lazy reader +read_lazy = dask.delayed(_read_image) +images = da.stack( + [ + da.from_delayed( + read_lazy(path), + shape=sample.shape, + dtype=sample.dtype, + ) + for path in img_paths + ] # list of dask arrays +) # (image_id, H, W, C), one chunk per image +print(images.shape) + +# add image array as data variable +ds["images"] = xr.DataArray( + data=images, + dims=["image_id", "img_h", "img_w", "rgb"], +) -# for one task per chunk: -def _apply_one_hot_encoding_for_IDs(block, n_ids): - """Apply one hot encoding to ID-encoded block. +# %%%%%%%%%%%%%%%%%%%%%% +# Get bbox, mask and ellipse of one box +image_id = 40 +id = 10 # annotation id +box = ds.sel(image_id=image_id, id=id) + +print(f"bbox centre (x, y): {box.position.values}") +print(f"bbox (width, height): {box.shape.values}") +print(f"mask centroid (x, y): {box.mask_centroid.values}") +print(f"mask area (pixels): {box.mask_area.values}") +print(f"ellipse major axis (pixels): {box.ellipse_major_axis_length.values}") +print( + "ellipse major axis theta (radians): " + f"{box.ellipse_major_axis_theta.values}" +) - 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) - # 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) - return block[:, None, :, :] == labels - - -# Apply expand labels chunk by chunk -# we define one chunk per image -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 - ), - dtype=bool, -) # (image_id, id, img_h, img_w) - -ds_bboxes["masks_bool"] = xr.DataArray( - data=mask_4d, - dims=["image_id", "id", "img_h", "img_w"], -) # %%%%%%%%%%%%%%%%%%%%%%% # Plot single mask and all masks in frame -# Show single mask for id=10 in frame 40 -image_id = 40 -id = 10 -single_mask = ds_bboxes.masks_bool.sel(image_id=image_id, id=id) +single_mask = ds.labels.sel(image_id=image_id) == box.label_of -fig, ax = plt.subplots(1, 1) # image -ax.imshow(ds_bboxes.image_array[image_id]) +fig, ax = plt.subplots(1, 1) +ax.imshow(ds.images.sel(image_id=image_id)) # bbox centre ax.scatter( - ds_bboxes.position.sel(image_id=image_id, id=id, space="x"), - ds_bboxes.position.sel(image_id=image_id, id=id, space="y"), + ds.position.sel(image_id=image_id, id=id, space="x"), + ds.position.sel(image_id=image_id, id=id, space="y"), 15, marker="x", color="r", @@ -219,10 +225,24 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids): cmap="Reds", alpha=single_mask.astype(float) * 0.5, ) + +# add contour ax.contour(single_mask, levels=[0.5], colors="red", linewidths=0.5) +# add ellipse +ax.add_patch( + Ellipse( + xy=box.mask_centroid.values, # (x, y) + width=float(box.ellipse_major_axis_length), + height=float(box.ellipse_minor_axis_length), + angle=np.degrees(float(box.ellipse_major_axis_theta)), + fill=False, + edgecolor="yellow", + ) +) + # plot all masks in one frame (one color per ID) -label_frame = mask_da_array[image_id].compute() +label_frame = ds.labels.sel(image_id=image_id).values ax.imshow( label_frame, cmap="turbo", @@ -231,41 +251,24 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids): # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Visualise frames and masks in napari viewer -mask_array = mask_da_array viewer = napari.Viewer() # 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") + +# add frames +viewer.add_image(ds.images.data, name="frames", rgb=True) + +# add masks +mask_array = ds.labels.data +viewer.add_labels(mask_array, 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]): - # 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( - [np.full(corners.shape[0], fill_value=frame_idx), corners] - ) - all_ellipses.append(nd) - for axis in major_axes: - nd = np.column_stack( - [np.full(axis.shape[0], fill_value=frame_idx), axis] - ) - all_major_axes.append(nd) - for axis in minor_axes: - nd = np.column_stack( - [np.full(axis.shape[0], fill_value=frame_idx), axis] - ) - all_minor_axes.append(nd) +# add ellipses and axis lines as napari Shapes layers, +# from the stored ellipse parameters +all_ellipses, all_major_axes, all_minor_axes = ellipse_shapes_from_ds(ds) viewer.add_shapes( - all_ellipses, + list(all_ellipses), shape_type="ellipse", edge_color="yellow", edge_width=4, @@ -273,14 +276,14 @@ def _apply_one_hot_encoding_for_IDs(block, n_ids): name="ellipses", ) viewer.add_shapes( - all_minor_axes, + list(all_minor_axes), shape_type="line", edge_color="green", edge_width=4, name="minor axes", ) viewer.add_shapes( - all_major_axes, + list(all_major_axes), shape_type="line", edge_color="red", edge_width=4, diff --git a/scripts/generate_masks_from_bboxes.py b/scripts/generate_masks_from_bboxes.py index dd0abc61..03dd6bbe 100644 --- a/scripts/generate_masks_from_bboxes.py +++ b/scripts/generate_masks_from_bboxes.py @@ -1,7 +1,8 @@ """Generate SAM2 masks from provided bboxes annotations. Pass bounding boxes as prompts to SAM2, predict masks in those regions, -and save as an ID-encoded zarr store. +and save them as ID-encoded label images in a zarr store, together with +the bounding boxes and a few shape parameters per mask. This script is a one-off tool to run on annotated frames whose mask labels are only meaningful within a frame. @@ -13,19 +14,58 @@ annotations/ # contains COCO-format JSON VIA_JSON_combined_coco_gen.json -Output masks are written to: - /annotations/masks_.zarr - -The bbox format for SAM is a (N, 4) float array of -[x1, y1, x2, y2] in pixel coordinates, where x = column (width axis), -y = row (height axis), 1 is the xmin, ymin corner and 2 is the -xmax ymax corner. - -Usage (dependencies are auto-installed via uv): - uv run generate_masks_from_bboxes.py /path/to/data_dir - uv run generate_masks_from_bboxes.py /path/to/data_dir --batch-size 8 - uv run generate_masks_from_bboxes.py /path/to/data_dir \ +Output is written to an xarray dataset in a zarr store: + /annotations/bboxes_and_masks_.zarr + +with the following layout:: + + ├── image_id, space, id coordinates from ds_bboxes + ├── position, shape, category, ... (image_id, [space,] id) from ethology + ├── labels (image_id, img_h, img_w) uint16, + │ 0 = background + ├── label_of (id,) uint16, label = position + │ along id + 1 + ├── mask_centroid (image_id, space, id) float64 + └── mask_area, mask_area_to_convex_hull_area, + mask_area_to_mask_bbox_area, ellipse_major_axis_length, + ellipse_minor_axis_length, ellipse_major_axis_theta, + ellipse_eccentricity (image_id, id) float64 + +All variables share the ``image_id`` and ``id`` coordinates, so the mask +and ellipse parameters of a bounding box sit at the same (image_id, id) +as the box. Slots with no box, or with an empty mask, are NaN. + +The mask parameters are computed with scikit-image's +``regionprops_table`` on the stored ``labels``, so an occluded crab is +measured on its visible pixels only ("smallest_wins" occlusion policy). +The ellipse is the one with the same second moments as the mask: + +- ``mask_centroid``: mean (x, y) of the mask pixels, and ellipse centre. +- ``mask_area``: number of mask pixels. +- ``mask_area_to_convex_hull_area``: mask area / area of its convex hull. +- ``mask_area_to_mask_bbox_area``: mask area / area of the mask's own + axis-aligned bounding box (not the annotated bounding box). +- ``ellipse_major_axis_length``, ``ellipse_minor_axis_length``: full axis + lengths in pixels, not semi-axes. +- ``ellipse_major_axis_theta``: angle of the major axis from the image + x-axis towards the y-axis, in radians, in [0, pi). The image y-axis + points down, so a positive angle looks clockwise on screen. It can be + passed to matplotlib's ``Ellipse(..., angle=np.degrees(theta))`` + directly on an image axis. It is arbitrary for near-circular masks + (``ellipse_eccentricity`` close to 0). +- ``ellipse_eccentricity``: 0 for a circle, towards 1 for an elongated + ellipse. + +Usage +From the repo root: + uv run scripts/generate_masks_from_bboxes.py /path/to/data_dir + uv run scripts/generate_masks_from_bboxes.py /path/to/data_dir \ + --batch-size 8 + uv run scripts/generate_masks_from_bboxes.py /path/to/data_dir \ --model facebook/sam2.1-hiera-large + +Dependencies are auto-installed in an ephemeral environment via uv. +torch is installed from the CUDA 12.8 index (Linux/Windows). """ # /// script @@ -33,11 +73,14 @@ # dependencies = [ # "Pillow", # "zarr", +# "xarray", +# "dask", +# "scikit-image", # "torch>=2.5.1", # "torchvision>=0.20.1", # "sam-2 @ git+https://github.com/facebookresearch/sam2.git", # "huggingface_hub", -# "ethology", +# "ethology>=0.1.4", # after stable ID assignment (PR 164) # ] # # [tool.uv.sources] @@ -52,116 +95,258 @@ import argparse import sys +from collections.abc import Iterable from datetime import datetime from pathlib import Path +from typing import Any, cast +import dask.array as da import numpy as np import torch +import xarray as xr import zarr -from ethology.io.annotations import load_bboxes from PIL import Image -from sam2.sam2_image_predictor import SAM2ImagePredictor +from skimage.measure import regionprops_table + +# Map frame index to (N, 4) boxes and (N,) mask label IDs +# (label IDs are the mask label actual values) +BoxesAndLabelIDs = dict[int, tuple[np.ndarray, np.ndarray]] + +# Label of the mask of the box at position k along `id` is k + LABEL_OFFSET; +# label 0 is background +LABEL_OFFSET = 1 + +# Map regionprops_table property to data variable +PROPS_TO_VARS = { + "area": "mask_area", + "solidity": "mask_area_to_convex_hull_area", + "extent": "mask_area_to_mask_bbox_area", + "axis_major_length": "ellipse_major_axis_length", + "axis_minor_length": "ellipse_minor_axis_length", + "eccentricity": "ellipse_eccentricity", +} class ImageArrayLazy: """A lazy array for images in a list.""" - def __init__(self, img_paths): - """Initialise with a list of image file paths.""" - self.img_paths = sorted(img_paths) - # add image shape, assuming all have same as - # first sample - sample = np.array(Image.open(img_paths[0])) # H, W, C + def __init__(self, img_paths: Iterable[Path]) -> None: + """Initialise with a list of image file paths. + + The order of ``img_paths`` is preserved, so that index ``i`` + corresponds to the ``i``-th path in the list. + """ + self.img_paths = list(img_paths) + + # all images must have the same size, so that the masks fit in + # one array; check it from the file headers, before running SAM2 + sizes = {Image.open(p).size for p in self.img_paths} # (W, H) + if len(sizes) > 1: + raise ValueError(f"Images have different sizes (W, H): {sizes}") + + sample = np.array(Image.open(self.img_paths[0])) # H, W, C self.img_h, self.img_w, self.img_c = sample.shape - def __len__(self): + def __len__(self) -> int: """Return the number of images.""" return len(self.img_paths) - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> np.ndarray: """Load and return image at the given index as a numpy array.""" return np.array(Image.open(self.img_paths[idx])) @property - def shape(self): + def shape(self) -> tuple[int, int, int, int]: """Return the shape ``(n_images, height, width, channels)``.""" return (len(self.img_paths), self.img_h, self.img_w, self.img_c) # B, H, W, C -def get_bboxes_x1y1_x2y2_per_frame(ds_bboxes): - """Build a dict mapping frame index to (N, 4) bbox arrays. +def get_image_paths(ds_bboxes: xr.Dataset) -> list[Path]: + """Return image file paths in the order of the ``image_id`` coordinate. + + The ``i``-th path corresponds to the ``i``-th element of + ``ds_bboxes.image_id``. + """ + # Build full paths to images + images_dir = Path(ds_bboxes.attrs["images_directories"]) + img_paths = [ + images_dir / ds_bboxes.attrs["map_image_id_to_filename"][img_id] + for img_id in ds_bboxes.image_id.values + ] + + # Check if a file is missing + missing = [p for p in img_paths if not p.is_file()] + if missing: + raise FileNotFoundError( + f"{len(missing)} annotated images not found, e.g. {missing[0]}" + ) + return img_paths + + +def get_bboxes_and_label_ids_per_frame( + ds_bboxes: xr.Dataset, +) -> BoxesAndLabelIDs: + """Build a dict mapping frame index to ``boxes`` and ``label_ids``. + + The frame index is the position along the ``image_id`` coordinate. - Converts centre+shape format to [x1, y1, x2, y2] corners. + ``boxes`` is an (N, 4) float array of [x1, y1, x2, y2] in pixel + coordinates, as expected by SAM2: (x1, y1) is the top-left + (xmin, ymin) corner and (x2, y2) the bottom-right (xmax, ymax) + corner, with x = column and y = row. + + ``label_ids`` is an (N,) array with the label ID used to paint the + mask associated to each box in the label image. The label ID is the + box's position along the ``id`` dimension + 1 (0 is background). """ + # Get corners for all images x1y1 = ds_bboxes.position - ds_bboxes.shape / 2 x2y2 = ds_bboxes.position + ds_bboxes.shape / 2 - return { - idx: np.c_[ - x1y1.sel(image_id=idx).dropna(dim="id", how="all").values.T, - x2y2.sel(image_id=idx).dropna(dim="id", how="all").values.T, + # Loop thru images + boxes_and_label_ids: BoxesAndLabelIDs = {} + for idx in range(ds_bboxes.sizes["image_id"]): + # (n_ids, 4), with one row per id slot, NaN for empty slots + corners = np.c_[ + x1y1.isel(image_id=idx).transpose("id", "space").values, + x2y2.isel(image_id=idx).transpose("id", "space").values, ] - for idx in range(len(ds_bboxes.image_id)) - } + # select non-nan only + non_nan = ~np.isnan(corners).all(axis=1) + boxes_and_label_ids[idx] = ( + corners[non_nan], + np.flatnonzero(non_nan).astype(np.uint16) + LABEL_OFFSET, + # label ID = box's position along `id` (before dropping NaNs) + # + LABEL_OFFSET + ) + return boxes_and_label_ids + +def create_dataset_store( + path_to_zarr: str | Path, + ds_bboxes: xr.Dataset, + image_shape: tuple[int, int], + metadata_dict: dict[str, Any], +) -> None: + """Initialise output zarr store. -def create_mask_zarr( - path_to_zarr, - zarr_array_shape, - metadata_dict=None, -): - """Create a zarr store for ID-encoded masks and write metadata.""" - n_images = zarr_array_shape[0] - image_h = zarr_array_shape[1] - image_w = zarr_array_shape[2] + Writes ds_bboxes, empty ``labels`` and ``label_of``. - mask_zarr = zarr.open( + ``labels`` is an (image_id, img_h, img_w) uint16 array of ID-encoded + masks, with 0 as background. + + ``label_of`` maps each ``id`` to the label of its mask, which is the + box's position along ``id`` + 1. + + The initialised store is written with ``compute=False``: coordinates, + metadata and the bbox variables go to disk, but no ``labels`` chunk does. + Unwritten chunks read back as the fill value, 0. + """ + # Get dimension sizes + image_h, image_w = image_shape + n_images, n_ids = ds_bboxes.sizes["image_id"], ds_bboxes.sizes["id"] + # the largest label is n_ids - 1 + LABEL_OFFSET + if n_ids - 1 + LABEL_OFFSET > np.iinfo(np.uint16).max: + raise ValueError( + f"{n_ids} id slots do not fit in uint16 labels " + f"(max {np.iinfo(np.uint16).max})" + ) + + # Define dataset structure + # zarr 3 has no stable format for numpy's fixed-width strings, so store + # the `space` coordinate as an object. ethology >= 0.1.4 + # already returns it as object; the cast covers datasets built + # before then or 'manually', where ["x", "y"] becomes dict[str, Any]: + """Return a copy of ``attrs`` with every ``Path`` turned into ``str``. + + zarr cannot store ``Path`` attributes. Dict keys are left unchanged. + """ + + def fix(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, dict): + return {k: fix(v) for k, v in value.items()} + if isinstance(value, list | tuple): + return [fix(v) for v in value] + return value + + return {key: fix(value) for key, value in attrs.items()} def predict_masks_across_images( - model, - image_array, - bbox_prompts_per_frame_idx, - output_zarr, - img_batch_size=4, -): - """Predict masks in batches and write to zarr store.""" + model: Any, + image_array: ImageArrayLazy, + boxes_and_label_ids_per_frame_idx: BoxesAndLabelIDs, + output_zarr: zarr.Array, + img_batch_size: int = 4, +) -> None: + """Predict masks in batches and write to zarr store. + + ``model`` is a ``sam2.sam2_image_predictor.SAM2ImagePredictor``. + """ + # Get dimensions n_frames = output_zarr.shape[0] # len(ds_bboxes.image_id) img_h, img_w = output_zarr.shape[ 1:3 ] # ds_bboxes.attrs["image_array"].shape[1:3] + # Loop in batches for batch_start in range(0, n_frames, img_batch_size): + # Get batch size actual_batch_size = min(img_batch_size, n_frames - batch_start) - # Integer array (B, H, W): - # 0 = background, 1+ = object instance IDs + # Get label arrays (B, H, W) + # label 0 = background label_mask_batch = np.zeros( - (actual_batch_size, img_h, img_w), dtype=np.int16 + (actual_batch_size, img_h, img_w), + dtype=np.uint16, ) - # List of (N_i, 4) arrays, one per image - boxes_batch = [ - bbox_prompts_per_frame_idx[f_i] - for f_i in range(batch_start, batch_start + actual_batch_size) - ] + # Get lists of (N_i, 4) boxes and (N_i,) label IDs, one per image + frame_idcs = range(batch_start, batch_start + actual_batch_size) + boxes_batch, label_ids_batch = zip( + *(boxes_and_label_ids_per_frame_idx[f_i] for f_i in frame_idcs), + strict=True, + ) # 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 = [] + idcs_rel_with_boxes: list[int] = [] for i, boxes in enumerate(boxes_batch): if len(boxes): idcs_rel_with_boxes.append(i) @@ -186,7 +371,6 @@ def predict_masks_across_images( ) # 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, @@ -198,14 +382,11 @@ def predict_masks_across_images( # 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) + label_mask_batch[idx_rel_batch] = flatten_masks_smallest_wins( + masks_one_frame, label_ids_batch[idx_rel_batch] + ) # log print( @@ -221,20 +402,153 @@ def predict_masks_across_images( ) # Mark frames as annotated in zarr store attributes - annotated = set(output_zarr.attrs.get("annotated_frames", [])) + annotated = set( + cast( + list[int], + output_zarr.attrs.get("annotated_frames", []), + ) + ) annotated.update(range(batch_start, batch_start + actual_batch_size)) output_zarr.attrs["annotated_frames"] = sorted(annotated) -def main(args): +def flatten_masks_smallest_wins( + masks: np.ndarray, + label_ids: np.ndarray, +) -> np.ndarray: + """Flatten (N, H, W) masks of one frame into an (H, W) label image. + + The pixels of mask ``i`` take the value ``label_ids[i]``, and 0 is + background. Masks are painted from largest to smallest area, so + where two masks overlap the smaller one keeps the shared pixels + ("smallest_wins" occlusion policy). + """ + masks = masks.astype(bool) + label_frame = np.zeros(masks.shape[1:], dtype=np.uint16) + areas = masks.reshape(len(masks), -1).sum(axis=1) + for i in np.argsort(-areas, kind="stable"): + label_frame[masks[i]] = label_ids[i] + return label_frame + + +def compute_mask_props_dataset( + labels_zarr: zarr.Array | np.ndarray, + ds_bboxes: xr.Dataset, +) -> xr.Dataset: + """Compute mask properties dataset. + + Masks are read from ``labels_zarr`` one frame at a time. Measuring the + stored masks, rather than the label images in memory during prediction, + is deliberate: the values can be reproduced from the store alone, and + this function can run on an existing store. + + The mask with label ``k + 1`` is measured and stored at position + ``k`` along ``id``, so each value sits at the (image_id, id) of its + bounding box. + + The ellipse is the one with the same second moments as the mask, as + defined by scikit-image. The returned dataset has these data variables, + with NaN for slots with no mask: + + - ``mask_centroid`` (image_id, space, id): mean (x, y) of the mask + pixels, which is also the ellipse centre. + - ``mask_area``: number of mask pixels. + - ``mask_area_to_convex_hull_area``: mask area / area of its convex + hull (scikit-image's ``solidity``). + - ``mask_area_to_mask_bbox_area``: mask area / area of the mask's own + axis-aligned bounding box (scikit-image's ``extent``). Note this is + the bounding box of the mask, not the annotated (prompt) bounding + box, so it depends on the mask's rotation. + - ``ellipse_major_axis_length``, ``ellipse_minor_axis_length``: full + axis lengths in pixels, not semi-axes. + - ``ellipse_major_axis_theta``: angle of the major axis from the image + x-axis (columns) towards the y-axis (rows), in radians, in [0, pi). + The image y-axis points down, so a positive angle looks clockwise + on screen. It is arbitrary for near-circular masks. + - ``ellipse_eccentricity``: 0 for a circle, towards 1 for an elongated + ellipse. + """ + # Get n images and max n annotations from ds_bboxes + n_images, n_ids = ds_bboxes.sizes["image_id"], ds_bboxes.sizes["id"] + + # Initialise one (image_id, id) array per scalar variable, and an + # (image_id, space, id) array for the centroid; NaN where no mask + scalar_var_names = [*PROPS_TO_VARS.values(), "ellipse_major_axis_theta"] + scalar_var_arrays = { + name: np.full((n_images, n_ids), np.nan) for name in scalar_var_names + } + centroid_array = np.full((n_images, 2, n_ids), np.nan) # 2 for space dim + + # Loop thru images + for f in range(n_images): + # Compute geom properties + table = regionprops_table( + np.asarray(labels_zarr[f]), + properties=["label", "centroid", "orientation", *PROPS_TO_VARS], + ) + + # Fill in scalar variable arrays + idx = table["label"] - LABEL_OFFSET + for prop_name, var_name in PROPS_TO_VARS.items(): + scalar_var_arrays[var_name][f, idx] = table[prop_name] + + # Convert geometric properties to the x-axis + # (scikit-image measures from the row axis) + scalar_var_arrays["ellipse_major_axis_theta"][f, idx] = np.mod( + np.pi / 2 - table["orientation"], np.pi + ) + centroid_array[f, 0, idx] = table["centroid-1"] # x = column + centroid_array[f, 1, idx] = table["centroid-0"] # y = row + + return xr.Dataset( + { + **{ + name: (("image_id", "id"), arr) + for name, arr in scalar_var_arrays.items() + }, + "mask_centroid": ( + ("image_id", "space", "id"), + centroid_array, + ), + }, + coords=ds_bboxes[["image_id", "space", "id"]].coords, + ) + + +def append_mask_props(path_to_zarr: str | Path, ds_props: xr.Dataset) -> None: + """Add the data variables of ``ds_props`` to an existing store. + + ``to_zarr`` writes values by position, without aligning them to the + stored coordinates, so ``ds_props`` must follow the store's + ``image_id`` and ``id`` order (it does when built by + ``compute_mask_props`` from the same ``ds_bboxes``). + """ + # drop coordinates from ds_props before appending, to not + # overwrite coordinates of the zarr store + ds_props = ds_props.drop_vars(list(ds_props.coords)) + + # copy attrs from store + # (fast, it only reads root metadata file) + ds_props.attrs = dict(zarr.open_group(path_to_zarr).attrs) + + # append ds_props to zarr store + # (to_zarr(mode="a") replaces the store's root attrs) + ds_props.to_zarr(path_to_zarr, mode="a") + + +def main(args: argparse.Namespace) -> None: """Generate masks from bbox annotations using SAM2.""" + # Imports not needed for tests + from ethology.io.annotations import load_bboxes + from sam2.sam2_image_predictor import SAM2ImagePredictor + # Parse CLI args data_dir = args.data_dir images_dir = data_dir / "frames" annotations_dir = data_dir / "annotations" annotations_file = annotations_dir / "VIA_JSON_combined_coco_gen.json" batch_size = args.batch_size - model_id = args.model + model_name = args.model # Set device if torch.cuda.is_available(): @@ -245,69 +559,80 @@ def main(args): device = "cpu" print(f"Using device: {device}") - # Create output zarr store + # Define output zarr store, holding bboxes, masks and mask parameters timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - output_masks_zarr = annotations_dir / f"masks_{timestamp}.zarr" - output_masks_zarr.mkdir(parents=True, exist_ok=True) + output_zarr = annotations_dir / f"bboxes_and_masks_{timestamp}.zarr" - # Read groundtruth as ethology annotation dataset + # Read groundtruth as an ethology annotation dataset ds_bboxes = load_bboxes.from_files( annotations_file, format="COCO", images_dirs=images_dir, ) - # Load ground truth images as lazy array - png_files = sorted(ds_bboxes.attrs["images_directories"].glob("*.png")) - image_array = ImageArrayLazy(png_files) - # ds_bboxes.attrs["image_array"] = image_array + # Load ground truth images as a lazy array + # (aligned with the image_id coordinate in ds_bboxes) + img_paths = get_image_paths(ds_bboxes) + image_array = ImageArrayLazy(img_paths) # Load SAM2 predictor on device image_predictor = SAM2ImagePredictor.from_pretrained( - model_id=model_id, + model_id=model_name, device=device, ) - # Create zarr store for output ID-encoded masks - n_images, image_h, image_w = image_array.shape[:3] # (T, H, W) + # Create zarr store with the bboxes, and an empty array + # for the output ID-encoded masks + image_h, image_w = image_array.shape[1:3] metadata_dict = { "timestamp": timestamp, - "sam2_model": model_id, + "sam2_model": model_name, "source_images_dir": str(images_dir), "annotations_file": str(annotations_file), "annotations_format": "COCO", - "n_images": n_images, - "image_shape": [image_h, image_w], "batch_size": batch_size, "multimask_output": False, "prompt_type": "bounding_box", "mask_encoding": "instance_id", "background_label": 0, - "id_offset": 1, + "id_offset": LABEL_OFFSET, + "occlusion_policy": "smallest_wins", "device": device, } - mask_zarr = create_mask_zarr( - output_masks_zarr, - (n_images, image_h, image_w), + create_dataset_store( + output_zarr, + ds_bboxes, + (image_h, image_w), metadata_dict=metadata_dict, ) - # Compute bbox prompts per frame, - # as dict mapping frame index to + # Get labels zarr array + # (one frame per chunk) + labels_zarr = zarr.open_array(output_zarr, path="labels") + + # Extract bbox prompts and label IDs per frame, + # as dict mapping frame index to a tuple of # (N, 4) bbox arrays in [x1, y1, x2, y2] format - map_frame_idx_to_boxes = get_bboxes_x1y1_x2y2_per_frame(ds_bboxes) + # and (N,) arrays with the mask label ID associated to each box + map_frame_idx_to_boxes_and_label_ids = get_bboxes_and_label_ids_per_frame( + ds_bboxes + ) # Predict masks in image batches # and write to zarr store in OCTRON-format predict_masks_across_images( image_predictor, image_array, - map_frame_idx_to_boxes, - mask_zarr, + map_frame_idx_to_boxes_and_label_ids, + labels_zarr, batch_size, ) - print(f"Saved ID-encoded mask zarr to {output_masks_zarr}") + # Add geom parameters per bbox to zarr store + ds_props = compute_mask_props_dataset(labels_zarr, ds_bboxes) + append_mask_props(output_zarr, ds_props) + + print(f"Saved bboxes, masks and mask parameters to {output_zarr}") def parse_args(list_args: list[str]) -> argparse.Namespace: diff --git a/tests/test_unit/test_generate_masks_from_bboxes.py b/tests/test_unit/test_generate_masks_from_bboxes.py new file mode 100644 index 00000000..57322c69 --- /dev/null +++ b/tests/test_unit/test_generate_masks_from_bboxes.py @@ -0,0 +1,497 @@ +from pathlib import Path +from typing import NamedTuple + +import numpy as np +import pytest +import xarray as xr +import zarr +from PIL import Image +from skimage.draw import ellipse + +from scripts.generate_masks_from_bboxes import ( + ImageArrayLazy, + append_mask_props, + compute_mask_props_dataset, + create_dataset_store, + flatten_masks_smallest_wins, + get_bboxes_and_label_ids_per_frame, +) + + +@pytest.fixture +def ds_bboxes_annotations() -> xr.Dataset: + """Get a sample `ethology` annotations bounding box dataset. + + Two images, three id slots, as ethology >= 0.1.4 returns them. + + ethology's ``image_id`` is the image's index in the sorted list of all + images in the annotation file, so images without annotations leave gaps. + "image_id": [5, 7] mimics that, and make a mix-up between image ID and + position along ``image_id``. + """ + return xr.Dataset( + { + "position": ( + ("image_id", "space", "id"), + np.full((2, 2, 3), np.nan), + ), + "shape": ( + ("image_id", "space", "id"), + np.full((2, 2, 3), np.nan), + ), + "image_shape": ( + ("image_id", "space"), + np.full((2, 2), [120, 100]), + ), + "category": (("image_id", "id"), np.full((2, 3), -1)), + }, + coords={ + "image_id": [5, 7], + "space": np.array(["x", "y"], dtype=object), + "id": range(3), + }, + attrs={ + "annotation_files": Path("/data/annotations/coco.json"), + "annotation_format": "COCO", + "images_directories": Path("/data/frames"), + "map_category_to_str": {1: "crab"}, + "map_image_id_to_filename": {5: "a.png", 7: "b.png"}, + }, + ) + + +class EllipseLabels(NamedTuple): + """Label image array and the ellipses drawn in them, in pixels.""" + + label_images: np.ndarray + centre_x: int + centre_y: int + semi_major_axis: int + semi_minor_axis: int + + +@pytest.fixture +def ellipse_labels(ds_bboxes_annotations) -> EllipseLabels: + """Get label image arrays with one elliptical mask per image. + + Mimics the ``labels`` array that ``predict_masks_across_images`` + writes to the zarr store for the two images in the + ``ds_bboxes_annotations`` fixture. Each pixel holds the label + of the mask it belongs to (label = id + 1), with 0 as background. + + The ellipses have known centroid, axis lengths and orientation, so + we can check the output of ``compute_mask_props_dataset``: + + - image_id 5 (index 0): id 0, major axis along x (theta = 0). + - image_id 7 (index 1): id 2, same ellipse rotated to theta = pi / 4. + """ + # Get shape of array from bboxes dataset + n_images = ds_bboxes_annotations.sizes["image_id"] + image_shape = ds_bboxes_annotations.image_shape.isel(image_id=0) + width = int(image_shape.sel(space="x")) + height = int(image_shape.sel(space="y")) + + # Define an ellipse centred in the image, with major axis along x + centre_x, centre_y = width // 2, height // 2 + semi_major_axis, semi_minor_axis = 20, 8 # pixels + + # Initialise label images array + label_images = np.zeros((n_images, height, width), np.uint16) + + # Define masks + # image_id 5: id 0 (label 1), unrotated + rows, cols = ellipse( + centre_y, # row + centre_x, # col + r_radius=semi_minor_axis, + c_radius=semi_major_axis, + ) + label_images[0, rows, cols] = 1 + + # image_id 7: id 2 (label 3), rotated. + # skimage.draw rotates from +x towards -y, so this is theta = pi / 4 + rows, cols = ellipse( + centre_y, # row + centre_x, # col + r_radius=semi_minor_axis, + c_radius=semi_major_axis, + rotation=-np.pi / 4, + ) + label_images[1, rows, cols] = 3 + + return EllipseLabels( + label_images, + centre_x, + centre_y, + semi_major_axis, + semi_minor_axis, + ) + + +def test_create_dataset_store(tmp_path, ds_bboxes_annotations): + """Check the initialised store is as expected. + + It should have zero labels, label_of and zarr-safe attrs. + """ + # Initialise store + n_images = ds_bboxes_annotations.sizes["image_id"] + image_h, image_w = 100, 120 + metadata = {"occlusion_policy": "smallest_wins"} + + zarr_store = tmp_path / "bboxes_and_masks.zarr" + create_dataset_store( + zarr_store, + ds_bboxes_annotations, + image_shape=(image_h, image_w), + metadata_dict=metadata, + ) + + # Inspect dataset from store + ds_init = xr.open_zarr(zarr_store) + + # check labels array + assert ds_init.labels.shape == (n_images, image_h, image_w) + # one frame per chunk on disk + assert ds_init.labels.encoding["chunks"] == (1, image_h, image_w) + + # check no labels chunk is written, so they read back as the fill value 0 + assert int(ds_init.labels.max()) == 0 + assert ds_init.labels.dtype == np.uint16 + + # check label_of maps each id slot to its label (id + 1; 0 is background) + assert list(ds_init.label_of.values) == [1, 2, 3] + + # check the bbox variables and coordinates are written unchanged, + # with image_id values kept (not reset to positions) + xr.testing.assert_equal( + ds_init[list(ds_bboxes_annotations.data_vars)], + ds_bboxes_annotations.drop_attrs(), + ) + + # check attributes are zarr-safe; int dict keys come back as strings + assert ds_init.attrs["images_directories"] == "/data/frames" + assert ds_init.attrs["map_image_id_to_filename"] == { + "5": "a.png", + "7": "b.png", + } + assert ds_init.attrs["occlusion_policy"] == metadata["occlusion_policy"] + + +def test_create_dataset_store_too_many_ids(tmp_path): + """Check an error is raised if the labels do not fit in uint16. + + Labels are id + 1 (0 is background), so the largest uint16 value + (65535) fits 65535 id slots. With one more, the last label would + wrap around to 0 and its mask would be lost as background. + """ + # Define a bboxes dataset with one id slot too many + n_ids = np.iinfo(np.uint16).max + 1 + ds_bboxes = xr.Dataset( + { + "position": ( + ("image_id", "id"), + np.full( + (1, n_ids), + np.nan, + ), + ), + }, + coords={ + "image_id": [0], + "space": ["x", "y"], + "id": range(n_ids), + }, + ) + + # Check that initialising the output zarr store with it + # raises an error + with pytest.raises(ValueError, match="do not fit in uint16"): + create_dataset_store( + tmp_path / "bboxes_and_masks.zarr", + ds_bboxes, + image_shape=(10, 10), + metadata_dict={}, + ) + + +@pytest.mark.parametrize( + "dims", + [ + ("image_id", "space", "id"), + ("image_id", "id", "space"), + ], +) +def test_get_bboxes_and_label_ids_per_frame(dims): + """Test output of get_bboxes_and_label_ids_per_frame. + + The output is a dict mapping position allong image_id to + boxes and label IDs. + + Boxes go from centre and size to [x1, y1, x2, y2] corners. Each + box's label ID is its id + 1. Frames are keyed by their position along + ``image_id`` (0, 1), not by their image_id value (5, 7). + + Parametrised over the dimension order, which must not change the + result. + """ + # Define two boxes in the first image, at id slots 0 and 2, + # as {id: (centre (x, y), size (width, height))} + # (id slot 1 and the second image have no boxes, so are NaN) + boxes_first_image = { + 0: ([10, 20], [4, 6]), + 2: ([50, 60], [10, 8]), + } + + # Build dataset with these boxes + position = np.full((2, 2, 3), np.nan) # (image_id, space, id) + shape = np.full((2, 2, 3), np.nan) + for id_slot, (centre, size) in boxes_first_image.items(): + position[0, :, id_slot] = centre + shape[0, :, id_slot] = size + + ds_bboxes = xr.Dataset( + { + "position": (("image_id", "space", "id"), position), + "shape": (("image_id", "space", "id"), shape), + }, + coords={ + "image_id": [5, 7], + "space": ["x", "y"], + "id": range(3), + }, + ).transpose(*dims) + + # Compute boxes and label IDs per frame + boxes_and_label_ids = get_bboxes_and_label_ids_per_frame(ds_bboxes) + + # Check keys are the position along image_id + assert list(boxes_and_label_ids.keys()) == [0, 1] + + # Check boxes in first frame + # corners as: + # [x1, y1, x2, y2] = [x - w/2, y - h/2, x + w/2, y + h/2], + boxes, label_ids = boxes_and_label_ids[0] + np.testing.assert_array_equal( + boxes, + [ + [8, 17, 12, 23], # id 0 + [45, 56, 55, 64], # id 2 + ], + ) + # with label IDs id + 1 + np.testing.assert_array_equal( + label_ids, + [id_slot + 1 for id_slot in boxes_first_image], + ) + + # check second frame has no boxes + boxes, label_ids = boxes_and_label_ids[1] + assert boxes.shape == (0, 4) + assert label_ids.shape == (0,) + + +def test_compute_mask_props(ellipse_labels, ds_bboxes_annotations): + """Check computation of mask properties are correct and aligned. + + They should land at the right (image_id, id) and be correct. + """ + # Compute mask properties dataset + ds_props = compute_mask_props_dataset( + ellipse_labels.label_images, + ds_bboxes_annotations, + ) + + # Select a single mask (first ellipse) + ds_single_mask = ds_props.sel(image_id=5, id=0) + + # check properties of first ellipse + # mask area + assert ( + ds_single_mask.mask_area == (ellipse_labels.label_images[0] == 1).sum() + ) + # check centroid, as (x, y) along "space" + np.testing.assert_allclose( + ds_single_mask.mask_centroid, + [ellipse_labels.centre_x, ellipse_labels.centre_y], + ) + # check axes lengths + np.testing.assert_allclose( + ds_single_mask.ellipse_major_axis_length, + 2 * ellipse_labels.semi_major_axis, + rtol=0.05, + ) + np.testing.assert_allclose( + ds_single_mask.ellipse_minor_axis_length, + 2 * ellipse_labels.semi_minor_axis, + rtol=0.1, + ) + # check orientation (theta) + np.testing.assert_allclose( + ds_single_mask.ellipse_major_axis_theta, + 0, + atol=0.02, + ) + # check eccentricity + np.testing.assert_allclose( + ds_single_mask.ellipse_eccentricity, + np.sqrt( + 1 + - ( + ds_single_mask.ellipse_minor_axis_length + / ds_single_mask.ellipse_major_axis_length + ) + ** 2 + ), + atol=1e-3, + ) + + # check theta in second ellipse + # Note: second ellipse is same as first but rotated + np.testing.assert_allclose( + ds_props.ellipse_major_axis_theta.sel(image_id=7, id=2), + np.pi / 4, + atol=0.02, + ) + + # check only the slots with a mask have values + has_mask = xr.zeros_like(ds_props.mask_area, dtype=bool) + has_mask.loc[{"image_id": 5, "id": 0}] = True + has_mask.loc[{"image_id": 7, "id": 2}] = True + for da in ds_props.data_vars.values(): + if "space" in da.dims: # e.g. mask_centroid + has_value = da.notnull().all("space") + else: + has_value = da.notnull() + xr.testing.assert_equal(has_value, has_mask) + + +@pytest.mark.parametrize( + "space_dtype", + [ + object, # as from ethology >= 0.1.4 + " 0), # > 0 turns 0s into NaNs + ds.mask_area.transpose(*pixel_counts.dims), + ) + + +def test_image_array_lazy_different_sizes(tmp_path): + """Check images of different sizes raise an error.""" + image_paths = [ + tmp_path / "a.png", + tmp_path / "b.png", + ] + # Create mock images at those paths with different sizes + Image.new("RGB", (120, 100)).save(image_paths[0]) + Image.new("RGB", (60, 100)).save(image_paths[1]) + + # Check that defining an image lazy array raises + with pytest.raises(ValueError, match="different sizes"): + ImageArrayLazy(iter(image_paths)) + + +def test_flatten_masks_smallest_wins(): + """Check overlapping pixels take the label of the smaller mask.""" + # Define an array with overlapping masks + masks = np.zeros((3, 10, 10), bool) + masks[0, 2:8, 2:8] = True # small, 36 pixels, wins in overlap + masks[1, :, :6] = True # large, 60 pixels, overlaps mask 0 + masks[2, 8:, 6:] = True # 8 pixels, no overlap + + # Define arbitrary label IDs (not 1, 2, 3), to check masks are + # painted with the given label IDs, not by their position + label_ids = [4, 7, 9] + + # Apply mask flattening (smallest wins) + label_frame = flatten_masks_smallest_wins(masks, label_ids) + + # check output type + assert label_frame.dtype == np.uint16 + + # check each mask region has its label ID + assert (label_frame[masks[0]] == label_ids[0]).all() + # mask 1 keeps its area minus the overlap with mask 0 + assert (label_frame[masks[1] & ~masks[0]] == label_ids[1]).all() + assert (label_frame[masks[2]] == label_ids[2]).all() + + # check background is 0 + assert (label_frame[~masks.any(axis=0)] == 0).all() + + +def test_flatten_masks_smallest_wins_equal_areas(): + """Check equal-area masks are painted in input order: last one wins.""" + # Define an array with overlapping masks of same area + masks = np.zeros((2, 10, 10), bool) + masks[0, :, :6] = True # columns 0-5, 60 pixels + masks[1, :, 4:] = True # columns 4-9, 60 pixels, overlaps mask 0 + + # Apply flattening + label_ids = [1, 2] + label_frame = flatten_masks_smallest_wins(masks, label_ids) + + # columns 4 and 5 overlap: they take the label ID of the last mask + expected = np.zeros((10, 10), np.uint16) + expected[:, :4] = label_ids[0] + expected[:, 4:] = label_ids[1] + np.testing.assert_array_equal(label_frame, expected)