[sfmig]

Description

What is this PR?

• Bug fix
• Addition of a new feature
• Other

Why is this PR needed?

scripts/generate_masks_from_bboxes.py predicts a SAM2 mask for every annotated bounding box, but it only saves the masks, as a bare zarr array. To use the masks downstream we need simple shape parameters per crab, such as area, centroid and the orientation of a fitted ellipse. Each parameter should be easy to trace back to the bounding box it came from. Until now the notebook recomputed the ellipses on the fly with regionprops, so they were never stored.

What does this PR do?

The script now writes a single xarray dataset to <data_dir>/annotations/bboxes_and_masks_<timestamp>.zarr. It holds the bounding boxes, the masks and the mask parameters, all on the same (image_id, id) grid:

├── 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 = id + 1
├── mask_centroid                     (image_id, space, id)     float64, NaN = no mask
└── 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, NaN = no mask

So everything about one box is a single .sel away:

box = ds.sel(image_id=40, id=10)
box.position, box.shape              # the bounding box
box.mask_centroid, box.mask_area     # its mask
box.ellipse_major_axis_theta         # its ellipse
ds.labels.sel(image_id=40) == box.label_of   # its mask pixels

Changes:

• create_dataset_store replaces create_mask_zarr.
  • It writes ds_bboxes, an empty lazy labels variable and label_of as a template, with to_zarr(..., compute=False), then returns the labels zarr array for SAM2 to fill. This is the same pattern as create_mask_store in Add mask-tracked-crabs entry point #293.
  • Path attributes are converted to str, because zarr cannot store them.
• compute_mask_props is new.
  • It reads labels back one frame at a time and measures each mask with scikit-image's regionprops_table.
  • Each value is stored at the (image_id, id) of its box: label k + 1 → id index k.
• Parameters (the ellipse is the one with the same second moments as the mask):
  • mask_centroid: mean (x, y) of the mask pixels, which is also the ellipse centre. This is not the bbox centre.
  • mask_area: number of mask pixels.
  • mask_area_to_convex_hull_area: scikit-image's solidity.
  • mask_area_to_mask_bbox_area: scikit-image's extent. It uses the mask's own axis-aligned box, not the annotated box, so it depends on rotation.
  • ellipse_major_axis_length and ellipse_minor_axis_length: full lengths, not semi-axes.
  • ellipse_major_axis_theta: angle from the image x-axis towards +y, in radians, in [0, π). It works directly with matplotlib: Ellipse(xy, w, h, angle=np.degrees(theta)).
  • ellipse_eccentricity.
• main:
  • casts the space coordinate to object, to avoid zarr 3's UnstableSpecificationWarning (fixed upstream in ethology#159);
  • appends the parameters to the store with append_mask_props, which keeps the store's root attributes (see below).
• Imports and dependencies: sam2 and ethology are now imported inside main, so that tests can import the script. xarray, dask and scikit-image are added to the inline script dependencies.
• Notebook: notebook_visualise_masks_from_zarr.py now reads the one store with xr.open_zarr, and draws the ellipses (matplotlib and napari) from the stored variables. ellipses_from_labels is removed.
• Earlier commits on this branch:
  • images follow the order of map_image_id_to_filename;
  • masks are saved as uint16;
  • mask labels come from the bbox id position;
  • overlapping masks follow the smallest wins occlusion policy.

Warning

Stores written before this PR (masks_<timestamp>.zarr) do not open in the updated notebook. Re-run the script to get a bboxes_and_masks_<timestamp>.zarr store, and set map_month_to_zarr in the notebook to its name.

Fix: root attributes were lost when appending the mask parameters

to_zarr(mode="a") replaces the store's root attributes with those of the appended dataset, and ds_props has none. So bboxes_and_masks_<timestamp>.zarr stores written before 09331ce have "attributes": {}. They lost:

• from ethology: annotation_files, annotation_format, images_directories, map_category_to_str and map_image_id_to_filename, which get_image_paths needs;
• from the script: timestamp, sam2_model, source_images_dir, annotations_file, annotations_format, batch_size, multimask_output, prompt_type, mask_encoding, background_label, id_offset, occlusion_policy and device.

annotated_frames survived, because it is on the labels array, not the root.

The new append_mask_props copies the stored attributes onto ds_props before appending. main and test_store_round_trip both use it.

Existing stores do not need SAM2 to run again. The script below rebuilds the root attributes from the COCO file with ethology and writes only the attributes. It stops if the store's images or boxes differ from the COCO file. sam2_model, batch_size and device were never saved anywhere else, so it fills in the CLI defaults and "unknown"; edit them if the run used other values. It also adds attrs_restored_on.

uv run restore_store_attrs.py <data_dir>/annotations/bboxes_and_masks_<timestamp>.zarr

restore_store_attrs.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["ethology==0.1.3", "xarray", "zarr>=3", "dask", "numpy"]
# ///
"""Restore the root attributes of a bboxes_and_masks_<timestamp>.zarr store.

Stores written before the fix in PR #298 lost their root attributes. This
rebuilds them from the COCO file the store was made from, without running
SAM2 again. Usage:

    uv run restore_store_attrs.py <data_dir>/annotations/bboxes_and_masks_<ts>.zarr
"""

import sys
from datetime import date
from pathlib import Path

import numpy as np
import xarray as xr
import zarr
from ethology.io.annotations import load_bboxes

store = Path(sys.argv[1])
data_dir = store.parent.parent
images_dir = data_dir / "frames"
annotations_file = data_dir / "annotations" / "VIA_JSON_combined_coco_gen.json"

ds_bboxes = load_bboxes.from_files(
    annotations_file, format="COCO", images_dirs=images_dir
)


def boxes_per_image(ds):
    """Return one sorted (n_boxes, 5) array per image: x, y, w, h, category.

    Sorted because the order of boxes along ``id`` may differ between
    ethology runs: it sorts with pandas' default, unstable, sort.
    """
    boxes = np.concatenate(
        [ds.position.values, ds.shape.values, ds.category.values[:, None]],
        axis=1,
    ).transpose(0, 2, 1)  # (image_id, id, 5)
    out = []
    for b in boxes:
        b = b[~np.isnan(b[:, 0])]
        out.append(b[np.lexsort(b.T[::-1])])
    return out


# the store must hold the same images and boxes as the COCO file,
# since the attributes describe the images by image_id
stored = xr.open_zarr(store)[["position", "shape", "category"]].load()
np.testing.assert_array_equal(stored.image_id, ds_bboxes.image_id)
for b_stored, b_coco in zip(
    boxes_per_image(stored), boxes_per_image(ds_bboxes), strict=True
):
    np.testing.assert_array_equal(b_stored, b_coco)

# as in main() of generate_masks_from_bboxes.py. sam2_model, batch_size and
# device were not saved: edit them if the run did not use these values
metadata_dict = {
    "timestamp": store.stem.removeprefix("bboxes_and_masks_"),
    "sam2_model": "facebook/sam2.1-hiera-base-plus",  # CLI default
    "source_images_dir": str(images_dir),
    "annotations_file": str(annotations_file),
    "annotations_format": "COCO",
    "batch_size": 4,  # CLI default
    "multimask_output": False,
    "prompt_type": "bounding_box",
    "mask_encoding": "instance_id",
    "background_label": 0,
    "id_offset": 1,
    "occlusion_policy": "smallest_wins",
    "device": "unknown",
    "attrs_restored_on": date.today().isoformat(),
}

# zarr cannot store Path attributes (as _attrs_for_zarr in the script)
attrs = {
    **ds_bboxes.attrs,
    "annotation_files": str(ds_bboxes.attrs["annotation_files"]),
    "images_directories": str(ds_bboxes.attrs["images_directories"]),
    **metadata_dict,
}
zarr.open_group(store).attrs.update(attrs)
zarr.consolidate_metadata(store)  # xr.open_zarr reads the consolidated copy

print(f"Restored {len(attrs)} root attributes in {store}")

Note

The store's boxes match a fresh ethology load per image, but in a different order along id, in all 544 images of sep2023-full. Two fresh loads in the same environment agree with each other. ethology sorts with pandas' default sort_values, which is not stable, so the order of boxes within an image may depend on the environment. This is not verified. Until it is, id should be read as meaningful only inside the store that defines it: downstream code should slice the store, not rebuild boxes from the COCO file. The restore script compares boxes regardless of order for this reason.

The design is in claude-plans/proposal-mask-ellipse-params.md. It covers the parameter list, why perimeter is left out, the angle convention, and why the parameters are measured on the stored labels.

References

• Add mask-tracked-crabs entry point #293 — Add mask-tracked-crabs entry point. This PR uses the same template + zarr-fill pattern and the same labels / label_of layout.
• Segment crabs with SAM2 after detection proposals #292 — Segment crabs with SAM2 after detection proposals. That PR leaves ellipse fitting and orientation out of scope.
• Store the space coordinate of bboxes datasets as object dtype neuroinformatics-unit/ethology#159 — upstream fix for the space coordinate dtype.

How has this PR been tested?

• New unit tests in tests/test_unit/test_generate_masks_from_bboxes.py. They use synthetic ellipses and need no SAM2 and no real data:
  • test_compute_mask_props: values land at the right (image_id, id); area, centroid, full axis lengths and θ (0 and π/4) are correct; empty slots are NaN.
  • test_create_dataset_store: the template has all-zero uint16 labels, label_of, the bbox variables and zarr-safe attributes.
  • test_store_round_trip: runs the three passes of main (template → fill labels → append parameters) and checks that the stored masks and areas agree, and that the root attributes survive the append. Without the fix this fails with KeyError: 'map_image_id_to_filename'.
• restore_store_attrs.py was run on copies of the local sep2023-full and aug2023-full stores (without the labels chunks). The checks pass, and 19 root attributes are written, including a map_image_id_to_filename with 544 entries for sep2023-full.
• The full unit test suite passes locally (135 passed).
• Manual checks on synthetic masks:
  • the notebook's new ellipse_shapes_from_ds gives the same napari ellipses and axes as the old regionprops-based ellipses_from_labels;
  • a matplotlib Ellipse drawn from the stored parameters covers 99.7 % of a rotated mask's pixels;
  • the annotated_frames attribute that predict_masks_across_images writes on labels survives xr.open_zarr.
• Not yet done: a run of the script with SAM2 on real data (e.g. aug2023-full), and a visual check in the notebook that the ellipses sit on the crabs. This would also confirm the attribute types that ethology sets.

Does this PR require an update to the documentation?

Yes, and it is updated here:

• the module docstring of scripts/generate_masks_from_bboxes.py now describes the store layout, every parameter and the angle convention;
• the notebook shows how to read the store and draw the stored ellipses.

Checklist:

• The code has been tested locally
• Tests have been added to cover all new functionality
• The documentation has been updated to reflect any changes
• The code has been formatted with pre-commit

Note

The check-manifest hook fails on this branch because claude-plans/proposal-mask-ellipse-params.md is tracked but not in the sdist. All other hooks pass. The proposal should probably be removed (or claude-plans/ added to the check-manifest ignore list) before merging.

🤖 Generated with Claude Code