Enable automatic segmentation state caching for SAM2

[anwai98]

WIP. Need to check how this works across tools!

[anwai98]

Okay I tried it on a Mac, looks pretty good. What do you think? @constantinpape

[constantinpape]

I would change / clarify a few things:

• Is the loading now also lazy / on-demand for tiled / 3d data?
• I would try to save the state in zarr, as part of the embedding data structure. Also add a respective attribute to the zarr so we can record the state of this.
• Rename amg_state to something more general.

[anwai98]

I'll answer the high-level questions here:

Is the loading now also lazy / on-demand for tiled / 3d data?

Yes, that's right -- it's lazy on-demand loading. I updated it a bit to compat with the catched zarr dispatcher now!

I would try to save the state in zarr, as part of the embedding data structure. Also add a respective attribute to the zarr so we can record the state of this.

Took care of it. Here's a high-level file structure inside the cached zarr:

embeddings.zarr/
    -> features/
        - ... (contains primary image embedding stuff -- an array / group when tiled)
    -> high_res_feats / pos_enc / fpn
        - ... sibling HieraViT/SAM2 features -- still in the embeddings family
    -> automatic_segmentation_state/
        - ais/
            - ... (contains segmentation decoder outputs as a chunked array)
        - amg/
            - ... (whole state pickeld to a byte array)

Rename amg_state to something more general.

This one's a bit of a thinker. I am thinking of going with a general alias, eg. autoseg_state? (let's talk quickly over lunch)

[anwai98]

Okay, this is ready for the next round of review @constantinpape