auroraomics 0.1.0.dev0__py3-none-any.whl

auroraomics/pack.py

@@ -0,0 +1,667 @@
+"""Build and validate the ``patches`` container.
+
+A prediction takes one ``.zip`` per sample: the tiles as images, a ``tiles.csv``
+naming each tile and where on the slide it came from, and optionally a
+thumbnail. It is a zip and not an array file because zip members stream — a
+loader can pull one tile at a time out of a 50k-tile archive, while a single
+packed array has to be materialised whole before the first tile is available.
+
+Two directions, with very different obligations:
+
+:func:`pack_tiles` **produces** an archive from tiles you already have. It runs
+    quality control as the tiles stream past, writes only the ones that pass,
+    strips every scrap of metadata from the images it writes, and reports what
+    it dropped and why. Metadata stripping is not cosmetic: a tile cut from a
+    scanned slide can inherit the slide label, the scanner serial and a capture
+    timestamp, which is patient-adjacent data leaving the building inside an
+    image nobody looks inside.
+
+:func:`read_archive` **consumes** one, and treats it as hostile. Everything it
+    can check from the archive's own index — member names, declared sizes, the
+    manifest, the image headers — is checked BEFORE a single pixel is decoded,
+    because a decoder is the largest attack surface in the whole path and the
+    sizes a zip declares are written by whoever produced it.
+"""
+
+from __future__ import annotations
+
+import csv
+import io
+import numbers
+import re
+import zipfile
+from collections import Counter
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from pathlib import Path, PurePosixPath
+from typing import Any
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from . import contracts
+from .qc import DEFAULT_THRESHOLDS, QcResult, QcThresholds, qc_tile
+
+#: The input kind whose registry row bounds an archive.
+KIND = "patches"
+
+#: Decimal places each QC statistic is written with. Percentages to two places
+#: and a variance to one is the precision the manifest carries, and a packer
+#: that writes more digits produces a file no committed expectation can pin.
+STAT_DECIMALS = {"foreground_ratio": 2, "hsv_ratio": 2, "laplacian_var": 1}
+
+#: Every member of a fixed archive timestamp, so two runs over the same tiles
+#: produce the same bytes. Zip cannot represent anything before 1980, and a
+#: real time would make the archive unhashable and every diff of two archives
+#: a diff of when they were built.
+_FIXED_TIMESTAMP = (1980, 1, 1, 0, 0, 0)
+
+#: Unix permissions for a written member: readable, never executable.
+_MEMBER_MODE = 0o644 << 16
+
+#: Image formats Pillow may report for a member, by extension.
+_FORMATS = {".png": {"PNG"}, ".jpg": {"JPEG"}, ".jpeg": {"JPEG"}}
+
+
+class PackError(ValueError):
+    """An archive cannot be built from what was passed."""
+
+
+class ArchiveError(ValueError):
+    """An archive is not a valid container and must not be trusted."""
+
+
+@dataclass(frozen=True)
+class Tile:
+    """One tile: an RGB ``uint8`` image and its position on the slide.
+
+    ``x`` and ``y`` are the tile's top-left corner in full-resolution slide
+    pixels, which is the frame every coordinate in the result is expressed in.
+    """
+
+    image: np.ndarray
+    x: int
+    y: int
+
+
+@dataclass(frozen=True)
+class Rejection:
+    """One tile that quality control dropped, and the numbers behind it."""
+
+    x: int
+    y: int
+    reason: str
+    stats: dict[str, float]
+
+
+@dataclass(frozen=True)
+class PackReport:
+    """What :func:`pack_tiles` wrote, and what it did not."""
+
+    path: Path
+    written: int
+    rejected: int
+    rejected_by_reason: dict[str, int]
+    rejections: tuple[Rejection, ...]
+    thresholds: QcThresholds
+    mpp: float
+
+    @property
+    def considered(self) -> int:
+        """Tiles handed in, whether they were written or not."""
+        return self.written + self.rejected
+
+
+@dataclass(frozen=True, slots=True)
+class ArchiveRow:
+    """One row of an archive's manifest, parsed.
+
+    ``slots=True`` and no untyped spillover dict: an archive may declare 50k
+    tiles and :class:`PatchArchive` holds every row for its whole open life, so
+    a per-row dict of the same values as strings is tens of megabytes of exact
+    duplicate. A reader that needs the extra QC columns should parse the
+    manifest itself rather than have every reader pay for them.
+    """
+
+    tile_id: str
+    file: str
+    x: int
+    y: int
+    width: int
+    height: int
+    mpp: float
+
+
+def _spec() -> dict[str, Any]:
+    return contracts.patch_archive()
+
+
+def _tile_name_pattern() -> re.Pattern[str]:
+    return re.compile(_spec()["tile_name_pattern"])
+
+
+def _name_for(index: int) -> str:
+    """The member name of the ``index``-th written tile.
+
+    The width of the number comes out of the container's own name pattern
+    rather than being spelled here, so the two cannot disagree — and when the
+    pattern cannot express an index, that is a hard error rather than a name a
+    validator on the other side will reject.
+    """
+    pattern = _spec()["tile_name_pattern"]
+    digits = re.search(r"\{(\d+)\}", pattern)
+    if digits is None:
+        # Falling back to a guessed width made the refusal below misstate the
+        # pattern's capacity as 10,000, which is a worse answer than saying
+        # the width could not be read.
+        raise PackError(
+            f"the container's tile name pattern {pattern} declares no digit "
+            "width, so a tile's name cannot be derived from it"
+        )
+    width = int(digits.group(1))
+    name = f"tile_{index:0{width}d}.png"
+    if not re.match(pattern, name):
+        raise PackError(
+            f"tile {index} cannot be named inside this container: the name "
+            f"pattern {pattern} holds at most {10 ** width} tiles. Reduce the "
+            "tile count first — auroraomics.subsample picks the densest square."
+        )
+    return name
+
+
+def _encode_png(image: np.ndarray) -> bytes:
+    """Encode RGB ``uint8`` as PNG carrying no metadata at all.
+
+    The encoder is handed nothing but pixels, so there is nothing for it to
+    write into a text chunk, and it writes no timestamp of its own. A test
+    walks the chunk types of a written tile to keep that true.
+    """
+    ok, buffer = cv2.imencode(".png", cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+    if not ok:  # pragma: no cover - only a broken build gets here
+        raise PackError("the image encoder refused a tile")
+    return buffer.tobytes()
+
+
+def _encode_jpeg(image: np.ndarray, quality: int = 90) -> bytes:
+    """Encode RGB ``uint8`` as JPEG carrying no EXIF and no thumbnail."""
+    ok, buffer = cv2.imencode(
+        ".jpg",
+        cv2.cvtColor(image, cv2.COLOR_RGB2BGR),
+        [int(cv2.IMWRITE_JPEG_QUALITY), quality],
+    )
+    if not ok:  # pragma: no cover
+        raise PackError("the image encoder refused the thumbnail")
+    return buffer.tobytes()
+
+
+def _write_member(archive: zipfile.ZipFile, name: str, payload: bytes) -> None:
+    info = zipfile.ZipInfo(name, date_time=_FIXED_TIMESTAMP)
+    info.compress_type = zipfile.ZIP_DEFLATED
+    info.external_attr = _MEMBER_MODE
+    archive.writestr(info, payload)
+
+
+def _format_stat(name: str, value: float) -> str:
+    return str(round(float(value), STAT_DECIMALS[name]))
+
+
+def pack_tiles(
+    tiles: Iterable[Tile],
+    out_zip: str | Path,
+    *,
+    mpp: float,
+    thresholds: QcThresholds | None = None,
+    thumbnail: np.ndarray | None = None,
+) -> PackReport:
+    """Write the tiles that pass quality control into one archive.
+
+    Parameters
+    ----------
+    tiles:
+        Any iterable of :class:`Tile`. It is consumed lazily and one tile at a
+        time, so a generator that decodes a slide region per step never holds
+        more than one tile — which is what makes a whole slide packable on a
+        laptop.
+    out_zip:
+        Where to write. An existing file is replaced.
+    mpp:
+        Micrometres per pixel of the tiles, at the resolution they were cut.
+        Recorded for every tile: without it a tile is a picture of unknown
+        physical size and no model can be matched to it.
+    thresholds:
+        Quality-control floors. The defaults if omitted.
+    thumbnail:
+        An optional RGB overview image, written as the container's thumbnail.
+        Any metadata it carries is dropped, as for every tile.
+
+    Returns
+    -------
+    PackReport
+        Counts, plus one :class:`Rejection` per dropped tile naming the
+        threshold that dropped it. A caller that gets an unexpectedly small
+        archive can see immediately whether the slide was empty or a threshold
+        was wrong.
+    """
+    # numbers.Real rather than (int, float), because a resolution usually comes
+    # out of a slide's own metadata as a numpy scalar; bool is excluded because
+    # True would otherwise be accepted as 1 micrometre per pixel.
+    if isinstance(mpp, bool) or not isinstance(mpp, numbers.Real) or not mpp > 0:
+        raise PackError(f"mpp must be a positive number of micrometres, got {mpp!r}")
+    thresholds = thresholds or DEFAULT_THRESHOLDS
+    spec = _spec()
+    columns = list(spec["manifest_columns"]) + list(QcResult.stat_names())
+
+    out_zip = Path(out_zip)
+    rows: list[list[str]] = []
+    rejections: list[Rejection] = []
+    by_reason: dict[str, int] = {}
+
+    with zipfile.ZipFile(out_zip, "w", zipfile.ZIP_DEFLATED) as archive:
+        for tile in tiles:
+            verdict = qc_tile(tile.image, thresholds)
+            if not verdict.passes:
+                reason = verdict.rejected_by or "unknown"
+                rejections.append(
+                    Rejection(
+                        x=int(tile.x), y=int(tile.y), reason=reason,
+                        stats=verdict.stats(),
+                    )
+                )
+                by_reason[reason] = by_reason.get(reason, 0) + 1
+                continue
+
+            name = _name_for(len(rows))
+            height, width = tile.image.shape[:2]
+            _write_member(archive, name, _encode_png(tile.image))
+            values = {
+                "tile_id": PurePosixPath(name).stem,
+                "file": name,
+                "x": str(int(tile.x)),
+                "y": str(int(tile.y)),
+                "width": str(int(width)),
+                "height": str(int(height)),
+                "mpp": str(mpp),
+                **{
+                    stat: _format_stat(stat, value)
+                    for stat, value in verdict.stats().items()
+                },
+            }
+            rows.append([values[column] for column in columns])
+
+        buffer = io.StringIO()
+        writer = csv.writer(buffer)
+        writer.writerow(columns)
+        writer.writerows(rows)
+        _write_member(archive, spec["manifest"], buffer.getvalue().encode("utf-8"))
+
+        if thumbnail is not None:
+            _write_member(
+                archive, spec["thumbnail"], _encode_jpeg(np.asarray(thumbnail))
+            )
+
+    return PackReport(
+        path=out_zip,
+        written=len(rows),
+        rejected=len(rejections),
+        rejected_by_reason=by_reason,
+        rejections=tuple(rejections),
+        thresholds=thresholds,
+        mpp=float(mpp),
+    )
+
+
+# ── reading someone else's archive ──────────────────────────────────────────
+
+
+class PatchArchive:
+    """A validated archive, open for reading one tile at a time.
+
+    Constructed by :func:`read_archive`, which has already checked everything
+    checkable without decoding. Decoding happens here, per tile, on demand.
+
+    The instance holds the archive OPEN, because a zip's index is read on open
+    and re-opening per tile would make reading an archive quadratic in the
+    number of tiles — which at 50k tiles is the difference between seconds and
+    an afternoon. Close it when done, or use it as a context manager:
+
+        >>> with read_archive("sample.zip") as archive:   # doctest: +SKIP
+        ...     for tile in archive.tiles():
+        ...         ...
+    """
+
+    def __init__(
+        self,
+        path: Path,
+        handle: zipfile.ZipFile,
+        rows: tuple[ArchiveRow, ...],
+        tile_names: tuple[str, ...],
+        has_thumbnail: bool,
+    ) -> None:
+        self.path = path
+        self.rows = rows
+        self.tile_names = tile_names
+        self.has_thumbnail = has_thumbnail
+        self._handle = handle
+        # read_tile checks membership once per tile and tiles() calls it once
+        # per row, so a linear scan of tile_names here is what makes reading a
+        # whole archive quadratic — the very cost this class exists to avoid.
+        self._tiles = frozenset(tile_names)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def close(self) -> None:
+        """Release the archive's file handle."""
+        self._handle.close()
+
+    def __enter__(self) -> PatchArchive:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    @property
+    def mpp(self) -> float:
+        """The archive's resolution, which every row must agree on."""
+        values = {row.mpp for row in self.rows}
+        if len(values) != 1:
+            raise ArchiveError(
+                f"the manifest declares {len(values)} different mpp values; "
+                "tiles of one sample are cut at one resolution"
+            )
+        return values.pop()
+
+    def _payload(self, name: str) -> bytes:
+        return _read_bounded(self._handle, name)
+
+    def read_tile(self, name: str) -> np.ndarray:
+        """Decode one tile to an RGB ``uint8`` array."""
+        if name not in self._tiles:
+            raise ArchiveError(f"{name!r} is not a tile of this archive")
+        image = cv2.imdecode(
+            np.frombuffer(self._payload(name), dtype=np.uint8), cv2.IMREAD_COLOR
+        )
+        if image is None:
+            raise ArchiveError(f"member {name} is not a decodable image")
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    def tiles(self) -> Iterator[Tile]:
+        """Every tile in manifest order, decoded one at a time."""
+        for row in self.rows:
+            yield Tile(image=self.read_tile(row.file), x=row.x, y=row.y)
+
+    def thumbnail(self) -> np.ndarray | None:
+        """The overview image as RGB ``uint8``, or ``None`` if there is none."""
+        if not self.has_thumbnail:
+            return None
+        payload = self._payload(_spec()["thumbnail"])
+        image = cv2.imdecode(np.frombuffer(payload, dtype=np.uint8), cv2.IMREAD_COLOR)
+        if image is None:
+            raise ArchiveError("the thumbnail is not a decodable image")
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+
+def _read_bounded(handle: zipfile.ZipFile, name: str) -> bytes:
+    """Read one member, refusing to accept more bytes than its index declares.
+
+    The declared size is what :func:`_check_declared_sizes` bounded, so it is
+    the number the caller has already agreed to spend memory on. Reading one
+    byte past it and refusing on a non-empty remainder turns a lying header
+    into an error instead of an unbounded read.
+    """
+    info = handle.getinfo(name)
+    with handle.open(info) as member:
+        payload = member.read(info.file_size + 1)
+    if len(payload) > info.file_size:
+        raise ArchiveError(
+            f"member {name} holds more bytes than its index declares — the "
+            "archive is inconsistent and was not decoded"
+        )
+    return payload
+
+
+def _check_member_name(name: str) -> None:
+    """Refuse any name that is not a plain file at the archive root.
+
+    A member name is attacker-controlled text, and the two classic shapes are
+    an absolute path and a ``../`` walk out of the extraction directory. This
+    package never extracts to disk, but a caller might, and a name that cannot
+    escape is a name nobody has to remember to sanitise.
+    """
+    if not name or name != PurePosixPath(name).name:
+        raise ArchiveError(
+            f"member {name!r} is not a plain file name at the archive root; a "
+            "container holds no directories and no paths"
+        )
+    if name.startswith(".") or "\\" in name or ".." in name:
+        raise ArchiveError(f"member {name!r} is not an acceptable file name")
+
+
+def _check_declared_sizes(infos: list[zipfile.ZipInfo], cap: int) -> None:
+    """Bound the archive by what it DECLARES, before anything is decompressed.
+
+    This is the zip-bomb check, and it has to happen here: a few hundred
+    kilobytes of deflate can declare — and expand to — gigabytes, so the only
+    safe moment to refuse is while the numbers are still just numbers.
+    """
+    total = 0
+    for info in infos:
+        if info.file_size < 0:
+            raise ArchiveError(f"member {info.filename} declares a negative size")
+        if info.file_size > cap:
+            raise ArchiveError(
+                f"member {info.filename} declares {info.file_size} bytes, over "
+                f"the {cap}-byte limit for one {KIND} input"
+            )
+        total += info.file_size
+    if total > cap:
+        raise ArchiveError(
+            f"the archive declares {total} bytes uncompressed, over the {cap}-byte "
+            f"limit for one {KIND} input"
+        )
+
+
+def _check_member_set(
+    names: list[str],
+    tile_pattern: re.Pattern[str],
+    manifest_name: str,
+    thumbnail_name: str,
+) -> tuple[str, ...]:
+    """Every member is named once, and named something the container defines.
+
+    A named step rather than two comprehensions inside :func:`read_archive`,
+    because it is the one part of validation that runs over EVERY member and so
+    is the one part whose cost has to be linear. Both checks were written as a
+    scan-inside-a-scan and cost ~49 s on a 50k-tile archive — the container's
+    own ceiling — for work that takes milliseconds. Keeping it here means it
+    can be tested at that ceiling without building a 50k-tile zip.
+
+    Returns the tile members, in archive order.
+    """
+    # A zip may physically hold two members with the SAME name, and which one a
+    # reader gets is up to the reader — so one tool validates a member and
+    # another decodes the other. There is no legitimate reason for a container
+    # to carry one, so it is a refusal rather than a choice.
+    duplicates = sorted(n for n, seen in Counter(names).items() if seen > 1)
+    if duplicates:
+        raise ArchiveError(
+            f"the archive holds more than one member named: "
+            f"{', '.join(duplicates)}"
+        )
+    tile_names = tuple(name for name in names if tile_pattern.match(name))
+    defined = frozenset(tile_names) | {manifest_name, thumbnail_name}
+    stray = [name for name in names if name not in defined]
+    if stray:
+        raise ArchiveError(
+            f"the archive holds member(s) the container does not define: "
+            f"{', '.join(sorted(stray))}"
+        )
+    return tile_names
+
+
+def _parse_manifest(payload: bytes, columns: list[str]) -> tuple[ArchiveRow, ...]:
+    try:
+        text = payload.decode("utf-8")
+    except UnicodeDecodeError as exc:
+        raise ArchiveError("the manifest is not UTF-8 text") from exc
+    reader = csv.reader(io.StringIO(text, newline=""))
+    try:
+        header = next(reader)
+    except StopIteration:
+        raise ArchiveError("the manifest is empty") from None
+    if header[: len(columns)] != columns:
+        raise ArchiveError(
+            f"the manifest must start with the columns {', '.join(columns)}; it "
+            f"starts with {', '.join(header[: len(columns)]) or '(nothing)'}"
+        )
+
+    rows: list[ArchiveRow] = []
+    for number, values in enumerate(reader, start=2):
+        if not values:
+            continue
+        if len(values) != len(header):
+            raise ArchiveError(
+                f"manifest line {number} has {len(values)} fields for "
+                f"{len(header)} columns"
+            )
+        raw = dict(zip(header, values))
+        try:
+            rows.append(
+                ArchiveRow(
+                    tile_id=raw["tile_id"],
+                    file=raw["file"],
+                    x=int(raw["x"]),
+                    y=int(raw["y"]),
+                    width=int(raw["width"]),
+                    height=int(raw["height"]),
+                    mpp=float(raw["mpp"]),
+                )
+            )
+        except ValueError as exc:
+            raise ArchiveError(f"manifest line {number}: {exc}") from exc
+    if not rows:
+        raise ArchiveError("the manifest declares no tiles")
+    return tuple(rows)
+
+
+def _check_image_header(name: str, payload: bytes) -> tuple[int, int]:
+    """Width and height from the image HEADER, without decoding the pixels.
+
+    Pillow parses the header and stops, and it refuses implausible dimensions
+    on its own, so the geometry can be checked while the pixel data is still
+    untouched. It also reports the real format, which catches a member whose
+    extension says one thing and whose bytes say another.
+    """
+    suffix = PurePosixPath(name).suffix.lower()
+    try:
+        with Image.open(io.BytesIO(payload)) as image:
+            width, height = image.size
+            found = image.format
+    except Exception as exc:  # Pillow raises a family of errors here
+        raise ArchiveError(f"member {name} is not a readable image: {exc}") from exc
+    if found not in _FORMATS.get(suffix, set()):
+        raise ArchiveError(
+            f"member {name} is a {found} file, which does not match its "
+            f"{suffix} extension"
+        )
+    if width <= 0 or height <= 0:
+        raise ArchiveError(f"member {name} declares an empty image")
+    return width, height
+
+
+def read_archive(zip_path: str | Path) -> PatchArchive:
+    """Validate an archive and return it, open for reading.
+
+    Every check a server has to make, in the order that keeps the risky ones
+    last:
+
+    1. it is a zip at all;
+    2. every member is a plain file name at the root — no paths, no traversal;
+    3. every DECLARED size, and their total, is inside the cap for this input
+       kind — before anything is decompressed;
+    4. the member names are exactly the manifest, the tiles matching the
+       container's name pattern, and at most the thumbnail: a stray member is a
+       refusal, not something to ignore;
+    5. the manifest parses, starts with the agreed columns, and its rows and
+       the tile members name each other exactly — in both directions;
+    6. every tile's header agrees with its manifest row and with every other
+       tile, so one sample is one geometry.
+
+    Raises :class:`ArchiveError` on the first failure, naming what failed.
+    """
+    zip_path = Path(zip_path)
+    spec = _spec()
+    manifest_name = spec["manifest"]
+    thumbnail_name = spec["thumbnail"]
+    tile_pattern = _tile_name_pattern()
+    columns = list(spec["manifest_columns"])
+
+    try:
+        archive = zipfile.ZipFile(zip_path)
+    except zipfile.BadZipFile as exc:
+        raise ArchiveError(f"{zip_path.name} is not a zip archive") from exc
+
+    # Every refusal below closes the handle on the way out: a rejected archive
+    # must not leave a file open in the caller's process.
+    try:
+        infos = [info for info in archive.infolist() if not info.is_dir()]
+        if not infos:
+            raise ArchiveError("the archive is empty")
+        for info in infos:
+            _check_member_name(info.filename)
+        _check_declared_sizes(infos, contracts.max_bytes(KIND))
+
+        names = [info.filename for info in infos]
+        tile_names = _check_member_set(
+            names, tile_pattern, manifest_name, thumbnail_name
+        )
+        if manifest_name not in names:
+            raise ArchiveError(f"the archive has no {manifest_name}")
+
+        rows = _parse_manifest(_read_bounded(archive, manifest_name), columns)
+
+        declared = [row.file for row in rows]
+        if len(set(declared)) != len(declared):
+            raise ArchiveError("the manifest names the same tile twice")
+        missing = sorted(set(declared) - set(tile_names))
+        if missing:
+            raise ArchiveError(
+                f"the manifest names tile(s) the archive does not contain: "
+                f"{', '.join(missing)}"
+            )
+        undeclared = sorted(set(tile_names) - set(declared))
+        if undeclared:
+            raise ArchiveError(
+                f"the archive contains tile(s) the manifest does not name: "
+                f"{', '.join(undeclared)}"
+            )
+
+        geometry: set[tuple[int, int]] = set()
+        for row in rows:
+            width, height = _check_image_header(
+                row.file, _read_bounded(archive, row.file)
+            )
+            if (width, height) != (row.width, row.height):
+                raise ArchiveError(
+                    f"member {row.file} is {width}x{height} but its manifest row "
+                    f"says {row.width}x{row.height}"
+                )
+            geometry.add((width, height))
+        if len(geometry) > 1:
+            sizes = ", ".join(f"{w}x{h}" for w, h in sorted(geometry))
+            raise ArchiveError(
+                f"the tiles are not one geometry ({sizes}); a model is matched to "
+                "one tile size and one resolution"
+            )
+    except BaseException:
+        archive.close()
+        raise
+
+    return PatchArchive(
+        path=zip_path,
+        handle=archive,
+        rows=rows,
+        tile_names=tile_names,
+        has_thumbnail=thumbnail_name in names,
+    )