From bdd148c6adc8f4b527191106b7ee2ecc7e0dc9d3 Mon Sep 17 00:00:00 2001
From: Talley Lambert <talley.lambert@gmail.com>
Date: Sat, 15 Jul 2023 15:23:51 -0400
Subject: [PATCH] fix: fix binary reshaping

---
 src/nd2/_binary.py                       | 12 ++++-----
 src/nd2/_sdk_types.py                    | 10 +++++++
 src/nd2/readers/_modern/modern_reader.py | 33 +++++++++++++-----------
 3 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/nd2/_binary.py b/src/nd2/_binary.py
index b2edb77..7a5431f 100644
--- a/src/nd2/_binary.py
+++ b/src/nd2/_binary.py
@@ -56,13 +56,13 @@ class BinaryLayer(NamedTuple):
 
     data: list[np.ndarray | None]
     name: str
-    comp_name: str
-    comp_order: int
-    color: int
-    color_mode: int
-    state: int
     file_tag: str
-    layer_id: int
+    comp_name: str | None
+    comp_order: int | None
+    color: int | None
+    color_mode: int | None
+    state: int | None
+    layer_id: int | None
     coordinate_shape: tuple[int, ...]
 
     @property
diff --git a/src/nd2/_sdk_types.py b/src/nd2/_sdk_types.py
index 32399f1..4f8ebd5 100644
--- a/src/nd2/_sdk_types.py
+++ b/src/nd2/_sdk_types.py
@@ -339,6 +339,16 @@ class RawTagDict(TypedDict):
         Desc: str
         Unit: str
 
+    class BinaryMetaDict(TypedDict):
+        BinLayerID: int
+        State: int
+        Color: int
+        CompOrder: int
+        Name: str
+        FileTag: str
+        CompName: str
+        ColorMode: int
+
     # These dicts are intermediate dicts created in the process of parsing raw meta
     # they mimic intermediate parsing done by the SDK... but needn't stay this way.
 
diff --git a/src/nd2/readers/_modern/modern_reader.py b/src/nd2/readers/_modern/modern_reader.py
index 8021a80..508abe0 100644
--- a/src/nd2/readers/_modern/modern_reader.py
+++ b/src/nd2/readers/_modern/modern_reader.py
@@ -38,6 +38,7 @@
     from nd2._binary import BinaryLayers
     from nd2._parse._chunk_decode import ChunkMap
     from nd2._sdk_types import (
+        BinaryMetaDict,
         GlobalMetadata,
         RawAttributesDict,
         RawExperimentDict,
@@ -541,7 +542,7 @@ def binary_data(self) -> BinaryLayers | None:
         binary_meta = self._decode_chunk(chunk_key, strip_prefix=True)
 
         try:
-            items: dict = binary_meta["BinaryMetadata_v1"]
+            items = cast("dict[str, BinaryMetaDict]", binary_meta["BinaryMetadata_v1"])
         except KeyError:  # pragma: no cover
             warnings.warn(
                 "Could not find 'BinaryMetadata_v1' tag, please open an "
@@ -550,28 +551,30 @@ def binary_data(self) -> BinaryLayers | None:
             )
             return None
 
-        binseqs = sorted(x for x in self.chunkmap if b"RleZipBinarySequence" in x)
         mask_items = []
         coord_shape = tuple(x.count for x in self.experiment())
-
         for _, item in sorted(items.items()):
-            key = item["FileTag"].encode()
+            # something like: RleZipBinarySequence_1bd900c
+            key = item["FileTag"]
             _masks: list[np.ndarray | None] = []
-            for bs in binseqs:
-                if key in bs:
-                    data = self._load_chunk(bs)[4:]
-                    _masks.append(decode_binary_mask(data) if data else None)
+            for plane in range(self._seq_count()):
+                # this will be something like
+                # b'CustomDataSeq|RleZipBinarySequence_1bd900c|1153!
+                chunk_key = f"CustomDataSeq|{key}|{plane}!".encode()
+                data = self._load_chunk(chunk_key)[4:]
+                _masks.append(decode_binary_mask(data) if data else None)
+
             mask_items.append(
                 BinaryLayer(
                     data=_masks,
-                    name=item["Name"],
-                    comp_name=item["CompName"],
-                    comp_order=item["CompOrder"],
-                    color_mode=item["ColorMode"],
-                    state=item["State"],
-                    color=item["Color"],
                     file_tag=key,
-                    layer_id=item["BinLayerID"],
+                    name=item["Name"],
+                    comp_name=item.get("CompName"),
+                    comp_order=item.get("CompOrder"),
+                    color_mode=item.get("ColorMode"),
+                    state=item.get("State"),
+                    color=item.get("Color"),
+                    layer_id=item.get("BinLayerID"),
                     coordinate_shape=coord_shape,
                 )
             )