feat: add roi parsing (#102)

* feat: add roi parsing

* test: update test link

* test: fix link

Author: tlambert03

README.md

@@ -66,6 +66,7 @@ f.attributes        # nd2.structures.Attributes
 f.metadata          # nd2.structures.Metadata
 f.frame_metadata(0) # nd2.structures.FrameMetadata (frame-specific meta)
 f.experiment        # List[nd2.structures.ExpLoop]
+f.rois              # Dict[int, nd2.structures.ROI]
 f.voxel_size()      # VoxelSize(x=0.65, y=0.65, z=1.0)
 f.text_info         # dict of misc info
 
@@ -230,6 +231,62 @@ Metadata(
 
 </details>
 
+
+<details>
+
+<summary><code>rois</code></summary>
+
+ROIs found in the metadata are available at `ND2File.rois`, which is a
+`dict` of `nd2.structures.ROI` objects, keyed by the ROI ID:
+
+```python
+{
+    1: ROI(
+        id=1,
+        info=RoiInfo(
+            shapeType=<RoiShapeType.Rectangle: 3>,
+            interpType=<InterpType.StimulationROI: 4>,
+            cookie=1,
+            color=255,
+            label='',
+            stimulationGroup=0,
+            scope=1,
+            appData=0,
+            multiFrame=False,
+            locked=False,
+            compCount=2,
+            bpc=16,
+            autodetected=False,
+            gradientStimulation=False,
+            gradientStimulationBitDepth=0,
+            gradientStimulationLo=0.0,
+            gradientStimulationHi=0.0
+        ),
+        guid='{87190352-9B32-46E4-8297-C46621C1E1EF}',
+        animParams=[
+            AnimParam(
+                timeMs=0.0,
+                enabled=1,
+                centerX=-0.4228425369685782,
+                centerY=-0.5194951478743071,
+                centerZ=0.0,
+                rotationZ=0.0,
+                boxShape=BoxShape(
+                    sizeX=0.21256931608133062,
+                    sizeY=0.21441774491682075,
+                    sizeZ=0.0
+                ),
+                extrudedShape=ExtrudedShape(sizeZ=0, basePoints=[])
+            )
+        ]
+    ),
+    ...
+}
+```
+
+</details>
+
+
 <details>
 
 <summary><code>text_info</code></summary>

scripts/download_samples.py

@@ -6,7 +6,7 @@
 import requests
 
 TEST_DATA = str(Path(__file__).parent.parent / "tests" / "data")
-URL = "https://www.dropbox.com/s/shbuvnkheudt7d7/nd2_test_data.zip?dl=1"
+URL = "https://www.dropbox.com/s/q57orjfzzagzull/nd2_test_data.zip?dl=1"
 
 
 def main():
@@ -23,7 +23,7 @@ def main():
             dl += len(data)
             f.write(data)
             done = int(50 * dl / total_length)
-            sys.stdout.write("\r[{}{}]".format("=" * done, " " * (50 - done)))
+            sys.stdout.write(f'\r[{"=" * done}{" " * (50 - done)}]')
             sys.stdout.flush()
     with ZipFile(f) as zf:
         zf.extractall(TEST_DATA)

src/nd2/nd2file.py

@@ -23,7 +23,7 @@
 import numpy as np
 
 from ._util import AXIS, VoxelSize, get_reader, is_supported_file
-from .structures import Attributes, ExpLoop, FrameMetadata, Metadata, XYPosLoop
+from .structures import ROI, Attributes, ExpLoop, FrameMetadata, Metadata, XYPosLoop
 
 try:
     from functools import cached_property
@@ -44,6 +44,9 @@
 
 Index = Union[int, slice]
 
+ROI_METADATA = "CustomData|RoiMetadata_v1"
+IMG_METADATA = "ImageMetadataLV"
+
 
 class ReadMode(str, Enum):
     MMAP = "mmap"
@@ -155,6 +158,21 @@ def text_info(self) -> Dict[str, Any]:
         """Misc text info."""
         return self._rdr.text_info()
 
+    @cached_property
+    def rois(self) -> Dict[int, ROI]:
+        """Return dict of {id: ROI} for all ROIs found in the metadata."""
+        if self.is_legacy or ROI_METADATA not in self._rdr._meta_map:  # type: ignore
+            return {}
+        data = self.unstructured_metadata(include={ROI_METADATA})
+        data = data.get(ROI_METADATA, {}).get("RoiMetadata_v1", {})
+        data.pop("Global_Size", None)
+        try:
+            _rois = (ROI._from_meta_dict(d) for d in data.values())
+            rois = {r.id: r for r in _rois}
+        except Exception as e:
+            raise ValueError(f"Could not parse ROI metadata: {e}") from e
+        return rois
+
     @cached_property
     def experiment(self) -> List[ExpLoop]:
         """Loop information for each nd axis"""
@@ -164,16 +182,16 @@ def experiment(self) -> List[ExpLoop]:
         # the SDK doesn't always do a good job of pulling position names from metadata
         # here, we try to extract it manually.  Might be error prone, so currently
         # we just ignore errors.
-        if not self.is_legacy and "ImageMetadataLV" in self._rdr._meta_map:  # type: ignore # noqa
+        if not self.is_legacy and IMG_METADATA in self._rdr._meta_map:  # type: ignore
             for n, item in enumerate(exp):
                 if isinstance(item, XYPosLoop):
                     names = {
                         tuple(p.stagePositionUm): p.name for p in item.parameters.points
                     }
                     if not any(names.values()):
                         _exp = self.unstructured_metadata(
-                            include={"ImageMetadataLV"}, unnest=True
-                        )["ImageMetadataLV"]
+                            include={IMG_METADATA}, unnest=True
+                        )[IMG_METADATA]
                         if n >= len(_exp):
                             continue
                         with contextlib.suppress(Exception):
@@ -238,7 +256,7 @@ def unstructured_metadata(
             _keys: Set[str] = set()
             for i in include:
                 if i not in keys:
-                    warnings.warn(f"include key {i!r} not found in metadata")
+                    warnings.warn(f"Key {i!r} not found in metadata")
                 else:
                     _keys.add(i)
             keys = _keys
@@ -253,7 +271,7 @@ def unstructured_metadata(
                     decoded: Any = meta.decode("utf-8")
                 else:
                     decoded = decode_metadata(meta, strip_prefix=strip_prefix)
-                    if key == "ImageMetadataLV" and unnest:
+                    if key == IMG_METADATA and unnest:
                         decoded = unnest_experiments(decoded)
             except Exception:
                 decoded = meta

src/nd2/structures.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import builtins
 from dataclasses import dataclass, field
 from enum import Enum, IntEnum
 from typing import List, NamedTuple, Optional, Tuple, Union
@@ -320,3 +321,153 @@ class Coordinate(NamedTuple):
     index: int
     type: str
     size: int
+
+
+def _lower0(x: str) -> str:
+    return x[0].lower() + x[1:]
+
+
+class BoxShape(NamedTuple):
+    sizeX: float = 0
+    sizeY: float = 0
+    sizeZ: float = 0
+
+
+class XYPoint(NamedTuple):
+    x: float = 0
+    y: float = 0
+
+
+class XYZPoint(NamedTuple):
+    x: float = 0
+    y: float = 0
+    z: float = 0
+
+
+class ExtrudedShape(NamedTuple):
+    sizeZ: float = 0
+    basePoints: List[XYPoint] = []
+
+    @classmethod
+    def _from_meta_dict(cls, val: dict) -> ExtrudedShape:
+        return cls(
+            sizeZ=val.get("SizeZ") or val.get("sizeZ") or 0,
+            basePoints=[
+                XYPoint(*val[f"BasePoints_{i}"].get("", []))
+                for i in range(val.get("BasePoints_Size", 0))
+            ],
+        )
+
+
+@dataclass
+class ROI:
+    """ROI object from NIS Elements."""
+
+    id: int
+    info: RoiInfo
+    guid: str
+    animParams: List[AnimParam] = field(default_factory=list)
+
+    def __post_init__(self):
+        self.info = RoiInfo(**self.info)
+        self.animParams = [AnimParam(**i) for i in self.animParams]
+
+    @classmethod
+    def _from_meta_dict(cls, val: dict) -> ROI:
+        anim_params = [
+            {_lower0(k): v for k, v in val[f"AnimParams_{i}"].items()}
+            for i in range(val.pop("AnimParams_Size", 0))
+        ]
+        return cls(
+            id=val["Id"],
+            info={_lower0(k): v for k, v in val["Info"].items()},
+            guid=val.get("GUID", ""),
+            animParams=anim_params,
+        )
+
+
+@dataclass
+class AnimParam:
+    """Parameters of ROI position/shape."""
+
+    timeMs: float = 0
+    enabled: bool = True
+    centerX: float = 0
+    centerY: float = 0
+    centerZ: float = 0
+    rotationZ: float = 0
+    boxShape: BoxShape = BoxShape()
+    extrudedShape: ExtrudedShape = ExtrudedShape()
+
+    def __post_init__(self):
+        if isinstance(self.boxShape, dict):
+            self.boxShape = BoxShape(
+                **{_lower0(k): v for k, v in self.boxShape.items()}
+            )
+        if isinstance(self.extrudedShape, dict):
+            self.extrudedShape = ExtrudedShape._from_meta_dict(self.extrudedShape)
+
+    @property
+    def center(self) -> XYZPoint:
+        """Center point as a named tuple (x, y, z)."""
+        return XYZPoint(self.centerX, self.centerY, self.centerZ)
+
+
+class RoiShapeType(IntEnum):
+    """The type of ROI shape."""
+
+    Raster = 1
+    Unknown2 = 2
+    Rectangle = 3
+    Ellipse = 4
+    Polygon = 5
+    Bezier = 6
+    Unknown7 = 7
+    Unknown8 = 8
+    Circle = 9
+    Square = 10
+
+
+class InterpType(IntEnum):
+    """The role that the ROI plays."""
+
+    StandardROI = 1
+    BackgroundROI = 2
+    ReferenceROI = 3
+    StimulationROI = 4
+
+
+@dataclass
+class RoiInfo:
+    """Info associated with an ROI."""
+
+    shapeType: RoiShapeType
+    interpType: InterpType
+    cookie: int = 0
+    color: int = 255
+    label: str = ""
+    # everything will default to zero, EVEN if "use as stimulation" is not checked
+    # use interpType to determine if it's a stimulation ROI
+    stimulationGroup: int = 0
+    scope: int = 1
+    appData: int = 0
+    multiFrame: bool = False
+    locked: bool = False
+    compCount: int = 2
+    bpc: int = 16
+    autodetected: bool = False
+    gradientStimulation: bool = False
+    gradientStimulationBitDepth: int = 0
+    gradientStimulationLo: float = 0.0
+    gradientStimulationHi: float = 0.0
+
+    def __post_init__(self):
+        # coerce types
+        for key, anno in self.__annotations__.items():
+            if key == "shapeType":
+                self.shapeType = RoiShapeType(self.shapeType)
+            elif key == "interpType":
+                self.interpType = InterpType(self.interpType)
+            else:
+                type_ = getattr(builtins, anno)
+                setattr(self, key, type_(getattr(self, key)))

tests/test_rois.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+import nd2
+
+DATA = Path(__file__).parent / "data"
+
+
+def test_rois():
+    with nd2.ND2File(DATA / "rois.nd2") as f:
+        rois = f.rois.values()
+        assert len(rois) == 18
+        assert [r.id for r in rois] == list(range(1, 19))