micro-reader 0.0.1__py3-none-any.whl

micro_reader/backends/mrc.py

@@ -0,0 +1,184 @@
+"""MRC (electron / light microscopy and cryo-EM volumes).
+
+Written from the public MRC2014 specification (CCP-EM): a 1024-byte header,
+an optional extended header of NSYMBT bytes, then raw pixels with X varying
+fastest, then Y, then Z (sections).  Uncompressed, so any rectangle is read
+exactly -- only the bytes it covers, straight into the result.  This
+replaces a JVM / Bio-Formats read in eubi-bridge.
+
+Supported: modes 0 (int8), 1 (int16), 2 (float32), 4 (complex64),
+6 (uint16), 12 (float16); little or big endian; extended headers; the
+standard axis order (MAPC, MAPR, MAPS) = (1, 2, 3).
+
+Orientation: MRC stores rows bottom-up (origin at the lower left, as IMOD
+displays it).  Images return them top-down, identical to Bio-Formats, so
+switching eubi-bridge between the two never flips an image.  (mrcfile and a
+raw memmap return file order: row 0 there is our last row.)
+Not yet: mode 3 (complex int16), mode 101 (4-bit), permuted axis orders.
+"""
+from __future__ import annotations
+
+import builtins
+import os
+import struct
+
+import numpy as np
+
+from .._base import File, Image
+from .._errors import CorruptFileError, UnsupportedFormatError
+from .._metadata import Channel, Metadata, quantity
+
+NAME = "mrc"
+HEADER_BYTES = 1024
+CHUNK_BYTES = 4 << 20        # full-width rows are read this many bytes at a time
+EXTENSIONS = {".mrc", ".mrcs", ".map", ".st", ".ali", ".rec"}
+
+#: MRC2014 mode -> numpy dtype.  Mode 0 is *signed* int8 in MRC2014; older
+#: writers used it unsigned, which cannot be told apart from the header.
+MODES = {0: "i1", 1: "i2", 2: "f4", 4: "c8", 6: "u2", 12: "f2"}
+UNSUPPORTED_MODES = {3: "complex int16", 101: "4-bit packed"}
+
+
+def sniff(head: bytes, path: str) -> bool:
+    """MRC2000+ files carry 'MAP ' at byte 208; older ones only the extension."""
+    if len(head) >= 212 and head[208:212] == b"MAP ":
+        return True
+    return os.path.splitext(path)[1].lower() in EXTENSIONS
+
+
+def _byte_order(head: bytes) -> str:
+    """'<' or '>' from the machine stamp (byte 212), else by plausibility."""
+    stamp = head[212] if len(head) > 212 else 0
+    if stamp in (0x44, 0x41):            # 0x44 0x44 (or 0x44 0x41): little endian
+        return "<"
+    if stamp == 0x11:                    # 0x11 0x11: big endian
+        return ">"
+    # No stamp (older writers): the order in which the mode is a known one
+    # and the dimensions are positive is the right one.
+    for order in ("<", ">"):
+        nx, ny, nz, mode = struct.unpack_from(order + "4i", head, 0)
+        if mode in MODES or mode in UNSUPPORTED_MODES:
+            if nx > 0 and ny > 0 and nz > 0:
+                return order
+    raise CorruptFileError("cannot determine the byte order of this MRC header")
+
+
+class _MrcImage(Image):
+    """The whole MRC volume as one image, axes 'zyx'."""
+
+    def __init__(self, path: str, offset: int, shape, file_dtype: np.dtype, sizes=None):
+        super().__init__(0, "zyx", shape, file_dtype, name=os.path.basename(path))
+        if sizes is not None:
+            self._metadata = Metadata(pixel_sizes=sizes, channels=(Channel(),))
+        self._path = path
+        self._offset = offset
+        self._file_dtype = file_dtype
+        self._closed = False
+
+    def _read_block(self, block):
+        # Plain reads straight into the result, one per row segment.  Not a
+        # memory map: mapped pages count as the process's memory (on Windows
+        # the whole file, after a full stream), which defeats the point of a
+        # memory-bounded reader.  Each call opens its own handle, so reads
+        # are thread-safe and nothing holds the file between reads.
+        if self._closed:
+            raise ValueError("this MRC image has been closed: its file was closed "
+                             "(File.close() or the end of a with block)")
+        zs, ys, xs = block
+        nz, ny, nx = self.shape
+        item = self._file_dtype.itemsize
+        out = np.empty((zs.stop - zs.start, ys.stop - ys.start, xs.stop - xs.start),
+                       dtype=self._file_dtype)
+        # MRC's origin is the lower-left corner: the first row stored is the
+        # bottom of the image.  Rows are returned top-down, as Bio-Formats
+        # (and so eubi-bridge, Fiji, OME-Zarr viewers) presents them: display
+        # row y is file row ny-1-y.
+        with builtins.open(self._path, "rb", buffering=0) as fh:
+            if xs.stop - xs.start == nx:
+                self._read_full_rows(fh, zs, ys, out)
+            else:
+                length = out.shape[2] * item
+                for i, z in enumerate(range(zs.start, zs.stop)):
+                    for j, y in enumerate(range(ys.start, ys.stop)):
+                        row = (z * ny + (ny - 1 - y)) * nx + xs.start
+                        fh.seek(self._offset + row * item)
+                        if fh.readinto(memoryview(out[i, j]).cast("B")) != length:
+                            raise CorruptFileError(
+                                f"{self._path!r}: short read at section {z}")
+        if not self._file_dtype.isnative:
+            out = out.byteswap(inplace=True).view(self.dtype)
+        return out, out.size
+
+    def _read_full_rows(self, fh, zs, ys, out):
+        """Full-width rows of a section are contiguous in the file: read them
+        in chunks of about CHUNK_BYTES and reverse each chunk into place.
+        Still exact; the only extra memory is one chunk."""
+        nz, ny, nx = self.shape
+        item = self._file_dtype.itemsize
+        per_chunk = max(1, CHUNK_BYTES // (nx * item))
+        scratch = np.empty((min(per_chunk, ys.stop - ys.start), nx), dtype=self._file_dtype)
+        first, stop = ny - ys.stop, ny - ys.start          # file rows, bottom-up
+        for i, z in enumerate(range(zs.start, zs.stop)):
+            for f in range(first, stop, per_chunk):
+                n = min(per_chunk, stop - f)
+                fh.seek(self._offset + (z * ny + f) * nx * item)
+                if fh.readinto(memoryview(scratch[:n]).cast("B")) != n * nx * item:
+                    raise CorruptFileError(f"{self._path!r}: short read at section {z}")
+                top = ny - f - n - ys.start                   # display row of file row f+n-1
+                out[i, top:top + n] = scratch[:n][::-1]
+
+    def close(self) -> None:
+        self._closed = True
+
+
+class _MrcFile(File):
+    backend = NAME
+
+
+def open(path: str, tiles: str = "separate", resolutions: str = "base") -> _MrcFile:  # noqa: A001 - backend protocol name
+    size = os.path.getsize(path)
+    if size < HEADER_BYTES:
+        raise CorruptFileError(f"{path!r} is smaller than an MRC header")
+    with builtins.open(path, "rb") as fh:
+        head = fh.read(HEADER_BYTES)
+
+    order = _byte_order(head)
+    nx, ny, nz, mode = struct.unpack_from(order + "4i", head, 0)
+    mapc, mapr, maps = struct.unpack_from(order + "3i", head, 64)
+    nsymbt, = struct.unpack_from(order + "i", head, 92)
+
+    if mode in UNSUPPORTED_MODES:
+        raise UnsupportedFormatError(
+            f"MRC mode {mode} ({UNSUPPORTED_MODES[mode]}) is not supported yet")
+    if mode not in MODES:
+        raise CorruptFileError(f"unknown MRC mode {mode} in {path!r}")
+    if min(nx, ny, nz) <= 0:
+        raise CorruptFileError(f"invalid MRC dimensions {(nx, ny, nz)} in {path!r}")
+    if (mapc, mapr, maps) not in ((1, 2, 3), (0, 0, 0)):
+        raise UnsupportedFormatError(
+            f"MRC axis order (MAPC, MAPR, MAPS) = {(mapc, mapr, maps)} is not "
+            "supported yet; only the standard (1, 2, 3)")
+    if nsymbt < 0:
+        raise CorruptFileError(f"negative extended header size {nsymbt} in {path!r}")
+
+    dtype = np.dtype(MODES[mode]).newbyteorder(order)
+    offset = HEADER_BYTES + nsymbt
+    needed = offset + nx * ny * nz * dtype.itemsize
+    if size < needed:
+        raise CorruptFileError(
+            f"{path!r} is truncated: the header describes {needed} bytes, "
+            f"the file has {size}")
+    return _MrcFile(path, [_MrcImage(path, offset, (nz, ny, nx), dtype, _pixel_sizes(head, order))])
+
+
+def _pixel_sizes(head: bytes, order: str) -> dict:
+    """Cell size (CELLA, Å) over the sampling (MX, MY, MZ): Å per pixel,
+    as MRC2014 defines it; kept in Å (OME-Zarr's ``angstrom``)."""
+    mx, my, mz = struct.unpack_from(order + "3i", head, 28)
+    cella = struct.unpack_from(order + "3f", head, 40)
+    sizes = {}
+    for axis, m, length in zip("xyz", (mx, my, mz), cella):
+        q = quantity(length / m, "angstrom") if m > 0 else None
+        if q:
+            sizes[axis] = q
+    return sizes

micro_reader/backends/nd2.py

@@ -0,0 +1,264 @@
+"""Nikon ND2 (version 3 and later; legacy versions 1 / 2 in ``_nd2_legacy``), own reader.
+
+A frame is one (time, position, Z) plane of all channels, stored in a chunk
+``ImageDataSeq|<n>!`` after an 8-byte timestamp: ``height`` rows of
+``width_bytes`` bytes (rows may be padded), each row ``width`` pixels of all
+channels and RGB components interleaved.  Frame numbers run over the
+experiment loops, outermost first.  Structure parsing is ported from nd2
+(Talley Lambert, BSD-3), see ``_nd2_structure``.
+
+- Uncompressed frames: only the requested rows are read (all channels
+  come along: they are interleaved per pixel).
+- "Lossless" frames are one zlib stream each: decoded whole, cached.
+
+Images are the XY positions, indexed by ``scene`` (as bioio-nd2 has them); axes ``tczyx``, plus ``s``
+for RGB components.  Legacy JPEG 2000 ND2 (versions 1 / 2): ``_nd2_legacy``.
+Not supported:
+lossy compression, tiled frames, custom loops.
+"""
+from __future__ import annotations
+
+import builtins
+import itertools
+import os
+import threading
+import zlib
+
+import numpy as np
+
+from .. import _memory, _pool, _stream
+from .._base import File, Image
+from .._errors import CorruptFileError, UnsupportedFormatError
+from .._io import DecodeCache, ThreadHandles
+from . import _nd2_legacy as legacy
+from . import _nd2_structure as structure
+
+NAME = "nd2"
+DECODE_CACHE_BYTES = 64 << 20
+RAW_TASK_BYTES = 8 << 20          # uncompressed frames are read in row pieces of this size
+STREAM_PIECE_BYTES = 8 << 20      # rows of a streamed "lossless" frame decoded at a time
+TIMESTAMP_BYTES = 8               # each frame chunk starts with a float64 timestamp
+
+
+def sniff(head: bytes, path: str) -> bool:
+    if head[:4] == structure.CHUNK_MAGIC.to_bytes(4, "little"):
+        return True
+    # legacy ND2 is a JPEG 2000 box file: the JP2 signature alone would also
+    # claim plain .jp2 images, so the extension decides (the box map at the
+    # end is checked on opening)
+    return legacy.is_legacy(head) and path.lower().endswith(".nd2")
+
+
+class Nd2Parser:
+    """What the images of one file share: layout, frame offsets, handles."""
+
+    def __init__(self, path: str):
+        self.path = path
+        with builtins.open(path, "rb") as fh:
+            ver = structure.version(fh.read(structure.START_CHUNK.size))
+            if ver < (3, 0):
+                raise UnsupportedFormatError(f"{path!r}: legacy ND2 version {ver}")
+            cmap = structure.chunkmap(fh)
+
+            def decoded(name):
+                if name not in cmap:
+                    return {}
+                return structure.decode_clx_lite(structure.read_chunk(fh, cmap[name][0]))
+
+            picture = decoded(b"ImageMetadataSeqLV|0!")
+            experiment = decoded(b"ImageMetadataLV!")
+            self.attrs = structure.attributes(decoded(b"ImageAttributesLV!"), picture)
+            self.loops = structure.experiment(experiment)
+            #: raw picture metadata and experiment, for ``_nd2_metadata``
+            self.picture = picture.get("SLxPictureMetadata", picture)
+            self.experiment = experiment
+        if self.attrs.compression == "lossy":
+            raise UnsupportedFormatError(f"{path!r}: lossy ND2 compression")
+        a = self.attrs
+        if a.components % a.channels:
+            raise UnsupportedFormatError(f"{path!r}: {a.components} components "
+                                         f"for {a.channels} channels")
+        self.samples = a.components // a.channels
+        self.dtype = np.dtype(a.dtype)
+        self.pixel_bytes = a.components * self.dtype.itemsize
+        if a.width_bytes < a.width * self.pixel_bytes:
+            raise CorruptFileError(f"{path!r}: rows of {a.width_bytes} bytes cannot hold "
+                                   f"{a.width} pixels of {self.pixel_bytes} bytes")
+        kinds = [loop.kind for loop in self.loops]
+        if len(set(kinds)) != len(kinds):
+            raise UnsupportedFormatError(f"{path!r}: repeated experiment loops {kinds}")
+        # frame chunk -> offset of its header; data offsets are read lazily
+        self.frames = {int(name[13:-1]): offset for name, (offset, _) in cmap.items()
+                       if name.startswith(b"ImageDataSeq|")}
+        self._data_at: dict = {}
+        self._lock = threading.Lock()
+        self.handles = ThreadHandles(path)
+
+    def size(self, kind: str) -> int:
+        return next((loop.count for loop in self.loops if loop.kind == kind), 1)
+
+    def frame_number(self, t: int, p: int, z: int) -> int:
+        index = {"T": t, "P": p, "Z": z}
+        number = 0
+        for loop in self.loops:
+            number = number * loop.count + index[loop.kind]
+        return number
+
+    def frame_data(self, number: int):
+        """(offset of the frame's pixels or compressed stream, its length),
+        or None for a frame that was never written."""
+        where = self._data_at.get(number)
+        if where is None:
+            header = self.frames.get(number)
+            if header is None:
+                return None
+            fh = self.handles.get()
+            start, length = structure.data_offset(fh, header)
+            where = self._data_at[number] = (start + TIMESTAMP_BYTES, length - TIMESTAMP_BYTES)
+        return where
+
+
+class _Nd2Image(Image):
+    """One XY position: axes 'tczyx' (+ 's' for RGB components)."""
+
+    def __init__(self, nd2: Nd2Parser, position: int):
+        a = nd2.attrs
+        self._f = nd2
+        self._p = position
+        self._rgb = nd2.samples > 1
+        shape = (nd2.size("T"), a.channels, nd2.size("Z"), a.height, a.width)
+        axes = "tczyx"
+        if self._rgb:
+            shape, axes = shape + (nd2.samples,), axes + "s"
+        #: decoded "lossless" frames kept for reuse
+        self.decode_cache = DecodeCache(DECODE_CACHE_BYTES)
+        self._closed = False
+        unit = (1, a.channels, 1, a.height, a.width) + ((nd2.samples,) if self._rgb else ()) \
+            if a.compression == "lossless" else None
+        super().__init__(position, axes, shape, nd2.dtype,
+                         name=f"{os.path.basename(nd2.path)} position {position}",
+                         read_unit=unit)
+
+    def _read_block(self, block):
+        if self._closed:
+            raise ValueError("this ND2 image has been closed: its file was closed "
+                             "(File.close() or the end of a with block)")
+        ts, cs, zs, ys, xs = block[:5]
+        ss = block[5] if self._rgb else slice(0, 1)
+        out = np.empty(tuple(s.stop - s.start for s in block), dtype=self.dtype)
+        view = out if self._rgb else out[..., None]          # always (..., s)
+        tasks = []
+        for t, z in itertools.product(range(ts.start, ts.stop), range(zs.start, zs.stop)):
+            dst = view[t - ts.start, :, z - zs.start]         # (c, y, x, s)
+            where = self._f.frame_data(self._f.frame_number(t, self._p, z))
+            if where is None:
+                dst[...] = 0
+            elif self._f.attrs.compression == "lossless":
+                tasks.append(lambda a=(t, z, where, dst, cs, ys, xs, ss): self._decoded_into(*a))
+            else:
+                tasks.extend(self._raw_tasks(where[0], dst, cs, ys, xs, ss))
+        return out, _pool.run(tasks)
+
+    def _frame_view(self, buf: np.ndarray, rows: int) -> np.ndarray:
+        """Bytes of *rows* whole rows -> (rows, x, channel, sample)."""
+        f, a = self._f, self._f.attrs
+        pixels = buf.reshape(rows, a.width_bytes)[:, :a.width * f.pixel_bytes]
+        return pixels.view(f.dtype).reshape(rows, a.width, a.channels, f.samples)
+
+    def _raw_tasks(self, start, dst, cs, ys, xs, ss):
+        f, a = self._f, self._f.attrs
+        per_task = max(1, RAW_TASK_BYTES // a.width_bytes)
+        whole_rows = (xs.stop - xs.start) * 2 >= a.width     # else: only the needed bytes per row
+        for y0 in range(ys.start, ys.stop, per_task):
+            y1 = min(ys.stop, y0 + per_task)
+            part = dst[:, y0 - ys.start:y1 - ys.start]
+            yield lambda args=(start, y0, y1, part, cs, xs, ss, whole_rows): self._read_rows(*args)
+
+    def _read_rows(self, start, y0, y1, dst, cs, xs, ss, whole_rows) -> int:
+        """Rows y0:y1 of an uncompressed frame into *dst* (c, y, x, s)."""
+        f, a = self._f, self._f.attrs
+        fh = f.handles.get()
+        n = y1 - y0
+        if whole_rows:
+            buf = np.empty(n * a.width_bytes, np.uint8)
+            fh.seek(start + y0 * a.width_bytes)
+            if fh.readinto(memoryview(buf)) != buf.nbytes:
+                raise CorruptFileError(f"{f.path!r}: frame data ends early")
+            pixels = self._frame_view(buf, n)[:, xs]
+            read = n * a.width * a.components
+        else:
+            width = xs.stop - xs.start
+            buf = np.empty((n, width * f.pixel_bytes), np.uint8)
+            for i in range(n):
+                fh.seek(start + (y0 + i) * a.width_bytes + xs.start * f.pixel_bytes)
+                if fh.readinto(memoryview(buf[i])) != buf.shape[1]:
+                    raise CorruptFileError(f"{f.path!r}: frame data ends early")
+            pixels = buf.view(f.dtype).reshape(n, width, a.channels, f.samples)
+            read = n * width * a.components
+        dst[...] = pixels[:, :, cs, ss].transpose(2, 0, 1, 3)
+        return read
+
+    def _decoded_into(self, t, z, where, dst, cs, ys, xs, ss) -> int:
+        frame, read = self.decode_cache.get((t, z)), 0
+        a = self._f.attrs
+        if frame is None and _memory.should_stream(a.height * a.width_bytes):
+            return self._streamed_into(where, dst, cs, ys, xs, ss)
+        if frame is None:
+            f, a = self._f, self._f.attrs
+            with _memory.reserve(a.height * a.width_bytes, f"{f.path!r}: an ND2 frame"):
+                data = zlib.decompress(f.handles.read(*where), 15, a.height * a.width_bytes)
+            if len(data) < a.height * a.width_bytes:
+                raise CorruptFileError(f"{f.path!r}: frame decodes to {len(data)} bytes")
+            frame = self._frame_view(np.frombuffer(data, np.uint8)[:a.height * a.width_bytes],
+                                     a.height)
+            self.decode_cache.put((t, z), frame)
+            read = frame.size
+        dst[...] = frame[ys, xs][:, :, cs, ss].transpose(2, 0, 1, 3)
+        return read
+
+    def _read_metadata(self):
+        from ._nd2_metadata import nd2_metadata
+        return nd2_metadata(self._f.picture, self._f.experiment, self._p,
+                            dict(zip(self.axes, self.shape)))
+
+    def _streamed_into(self, where, dst, cs, ys, xs, ss) -> int:
+        """Rows *ys* of a "lossless" frame too large to decode whole: its zlib
+        stream decoded as a stream, in pieces of rows (``_stream``)."""
+        f, a = self._f, self._f.attrs
+        per_piece = max(1, min(STREAM_PIECE_BYTES, _memory.memory_limit() // 4) // a.width_bytes)
+        for y0 in range(ys.start, ys.stop, per_piece):
+            y1 = min(ys.stop, y0 + per_piece)
+            with _memory.reserve((y1 - y0) * a.width_bytes, f"{f.path!r}: rows of a frame"):
+                buf = np.empty((y1 - y0) * a.width_bytes, np.uint8)
+                n = _stream.read(f.path, where[0], where[1], "zlib", y0 * a.width_bytes, buf,
+                                 size=a.height * a.width_bytes)
+                if n < buf.nbytes:
+                    raise CorruptFileError(f"{f.path!r}: frame ends at {n} of {buf.nbytes} bytes")
+                rows = self._frame_view(buf, y1 - y0)
+                dst[:, y0 - ys.start:y1 - ys.start] =                     rows[:, xs][:, :, cs, ss].transpose(2, 0, 1, 3)
+        return (ys.stop - ys.start) * a.width * a.components
+
+    def close(self) -> None:
+        self._closed = True
+        self.decode_cache.clear()
+        _stream.forget(self._f.path)
+
+
+class _Nd2File(File):
+    backend = NAME
+
+    def __init__(self, path, images, nd2: Nd2Parser):
+        super().__init__(path, images)
+        self._nd2 = nd2
+
+    def close(self) -> None:
+        super().close()
+        self._nd2.handles.close()
+
+
+def open(path: str, tiles: str = "separate", resolutions: str = "base"):  # noqa: A001 - backend protocol name
+    with builtins.open(path, "rb") as fh:
+        if legacy.is_legacy(fh.read(12)):
+            return legacy._LegacyNd2File(path)
+    nd2 = Nd2Parser(path)
+    return _Nd2File(path, [_Nd2Image(nd2, p) for p in range(nd2.size("P"))], nd2)