spatialscape 0.3.0__py3-none-any.whl

spatialscape/genes.py

@@ -0,0 +1,251 @@
+"""Gene models for the genome panel: one collapsed model per gene (the union of its exons), per chromosome.
+
+Sources (`tracks.genes` in the yaml): a GTF (Cell Ranger, GENCODE, Ensembl), a UCSC genePred table such as
+refGene.txt.gz, a BED (12 columns with blocks, or 3-6 columns of plain spans), or the keyword `refseq`, which
+downloads the UCSC RefSeq table for the genome once into the user cache (~/.cache/spatialscape). Written as
+
+    tracks/genes/index.json    {"source": ..., "chroms": {"chr1": ["DDX11L1", ...], ...}}   names, for search
+    tracks/genes/<chrom>.json  [[name, strand, start, end, [exon starts], [exon ends]], ...]  sorted by start
+
+Coordinates are 0-based half-open (BED style); strand is 1, -1, or 0 when unknown. The viewer loads the index
+once and one chromosome at a time.
+"""
+from __future__ import annotations
+
+import csv
+import json
+import os
+import shutil
+import urllib.request
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pandas as pd
+
+Log = Callable[[str], None]
+REFSEQ_URL = "https://hgdownload.soe.ucsc.edu/goldenPath/{genome}/database/refGene.txt.gz"
+REFSEQ_GENOMES = {"hg38", "hg19", "mm10", "mm39", "rn6", "rn7", "danRer11", "dm6", "ce11", "sacCer3"}  # UCSC assemblies with a refGene table
+KEYWORDS = {"refseq", "none"}
+STRAND = {"+": 1, "-": -1}
+# (chrom, name, strand, unit, exon start, exon end), 0-based half-open; `unit` tells transcripts (or rows) apart
+Record = tuple[str, str, int, object, int, int]
+
+
+@dataclass
+class GeneModel:
+    name: str
+    chrom: str
+    strand: int
+    start: int
+    end: int
+    exon_starts: list[int] = field(default_factory=list)
+    exon_ends: list[int] = field(default_factory=list)
+
+    def row(self) -> list:
+        return [self.name, self.strand, self.start, self.end, self.exon_starts, self.exon_ends]
+
+
+def cache_dir() -> Path:
+    root = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")
+    return Path(root) / "spatialscape"
+
+
+def normalize_chrom(chrom: str, sizes: Mapping[str, int]) -> str | None:
+    """Map `1`, `chr1`, `MT` onto the names in chrom.sizes; None drops the record (alt contigs, scaffolds)."""
+    if chrom in sizes:
+        return chrom
+    if chrom.startswith("chr"):
+        if chrom[3:] in sizes:
+            return chrom[3:]
+    elif f"chr{chrom}" in sizes:
+        return f"chr{chrom}"
+    if chrom in ("MT", "M", "chrMT", "chrM"):
+        return next((alt for alt in ("chrM", "MT", "M", "chrMT") if alt in sizes), None)
+    return None
+
+
+def merge_intervals(starts: Iterable[int], ends: Iterable[int]) -> tuple[list[int], list[int]]:
+    out_s: list[int] = []
+    out_e: list[int] = []
+    for s, e in sorted(zip(starts, ends, strict=True)):
+        if out_e and s <= out_e[-1]:
+            out_e[-1] = max(out_e[-1], e)
+        else:
+            out_s.append(s)
+            out_e.append(e)
+    return out_s, out_e
+
+
+def collapse(records: Iterable[Record]) -> dict[str, list[GeneModel]]:
+    """One model per locus of a gene name: transcripts (units) of the same name whose spans overlap are merged into
+    the union of their exons; copies elsewhere on the chromosome (the many U6 or Y_RNA genes) stay separate."""
+    acc: dict[tuple[str, str], dict[object, tuple[int, list[int], list[int]]]] = {}
+    for chrom, name, strand, unit, s, e in records:
+        units = acc.setdefault((chrom, name), {})
+        cur = units.get(unit)
+        if cur is None:
+            units[unit] = (strand, [s], [e])
+        else:
+            cur[1].append(s)
+            cur[2].append(e)
+    out: dict[str, list[GeneModel]] = {}
+    for (chrom, name), units in acc.items():
+        spans = sorted(((min(ss), max(ee), strand, ss, ee) for strand, ss, ee in units.values()), key=lambda t: (t[0], t[1]))
+        clusters: list[list[tuple[int, int, int, list[int], list[int]]]] = []
+        cluster_end = -1
+        for span in spans:
+            if clusters and span[0] <= cluster_end:
+                clusters[-1].append(span)
+                cluster_end = max(cluster_end, span[1])
+            else:
+                clusters.append([span])
+                cluster_end = span[1]
+        for cluster in clusters:
+            ms, me = merge_intervals([x for c in cluster for x in c[3]], [x for c in cluster for x in c[4]])
+            out.setdefault(chrom, []).append(GeneModel(name, chrom, cluster[0][2], ms[0], me[-1], ms, me))
+    for models in out.values():
+        models.sort(key=lambda g: (g.start, g.end, g.name))
+    return out
+
+
+def _chrom_map(chroms: Iterable[str], sizes: Mapping[str, int]) -> dict[str, str | None]:
+    return {c: normalize_chrom(c, sizes) for c in set(chroms)}
+
+
+def _read_gtf(path: Path, sizes: Mapping[str, int]) -> dict[str, list[GeneModel]]:
+    """Exon lines of a GTF, named by gene_name (gene_id when absent); streamed in chunks (GENCODE is ~3 M lines)."""
+    records: list[Record] = []
+    cols = ["chrom", "source", "feature", "start", "end", "score", "strand", "frame", "attrs"]
+    reader = pd.read_csv(
+        path, sep="\t", header=None, names=cols, usecols=["chrom", "feature", "start", "end", "strand", "attrs"],
+        dtype=str, quoting=csv.QUOTE_NONE, chunksize=500_000, na_filter=False,
+    )
+    for chunk in reader:
+        ex = chunk[chunk["feature"] == "exon"]
+        if ex.empty:
+            continue
+        cmap = _chrom_map(ex["chrom"].unique(), sizes)
+        chrom = ex["chrom"].map(cmap)
+        keep = chrom.notna()
+        if not keep.any():
+            continue
+        ex = ex[keep]
+        chrom = chrom[keep]
+        names = ex["attrs"].str.extract(r'gene_name "([^"]+)"', expand=False)
+        names = names.fillna(ex["attrs"].str.extract(r'gene_id "([^"]+)"', expand=False))
+        units = ex["attrs"].str.extract(r'transcript_id "([^"]+)"', expand=False).fillna(names)
+        ok = names.notna()
+        starts = ex["start"].astype("int64").to_numpy() - 1  # 1-based inclusive -> 0-based half-open
+        ends = ex["end"].astype("int64").to_numpy()
+        strands = ex["strand"].map(STRAND).fillna(0).astype(int).to_numpy()
+        for c, n, st, u, s, e, k in zip(chrom.to_numpy(), names.to_numpy(), strands, units.to_numpy(), starts, ends, ok.to_numpy(), strict=True):
+            if k:
+                records.append((c, n, int(st), u, int(s), int(e)))
+    return collapse(records)
+
+
+def _read_genepred(path: Path, sizes: Mapping[str, int]) -> dict[str, list[GeneModel]]:
+    """UCSC genePred / genePredExt, with or without the leading `bin` column of a table dump (refGene.txt.gz)."""
+    df = pd.read_csv(path, sep="\t", header=None, dtype=str, quoting=csv.QUOTE_NONE, na_filter=False, comment="#")
+    if df.shape[1] < 10:
+        raise ValueError(f"{path}: not a genePred table (expected at least 10 tab-separated columns)")
+    off = 1 if df.iloc[0, 3] in ("+", "-", ".") and df.shape[1] >= 11 else 0  # bin, name, chrom, strand, ...
+    name = df.iloc[:, off]
+    name2_col = off + 11
+    if df.shape[1] > name2_col:
+        name2 = df.iloc[:, name2_col]
+        name = name2.where(name2 != "", name)
+    chrom = df.iloc[:, off + 1].map(_chrom_map(df.iloc[:, off + 1].unique(), sizes))
+    strand = df.iloc[:, off + 2].map(STRAND).fillna(0).astype(int)
+    records: list[Record] = []
+    for i, (c, n, st, es, ee) in enumerate(zip(chrom, name, strand, df.iloc[:, off + 8], df.iloc[:, off + 9], strict=True)):
+        if c is None or not isinstance(c, str):
+            continue
+        for s, e in zip(es.strip(",").split(","), ee.strip(",").split(","), strict=True):
+            records.append((c, n, int(st), i, int(s), int(e)))
+    return collapse(records)
+
+
+def _read_bed(path: Path, sizes: Mapping[str, int]) -> dict[str, list[GeneModel]]:
+    """BED12 (blocks become exons) or plain BED3-6 spans (one block per gene, strand from column 6 when present)."""
+    df = pd.read_csv(path, sep="\t", header=None, dtype=str, quoting=csv.QUOTE_NONE, na_filter=False, comment="#")
+    ncol = df.shape[1]
+    chrom = df.iloc[:, 0].map(_chrom_map(df.iloc[:, 0].unique(), sizes))
+    records: list[Record] = []
+    for i, c in enumerate(chrom):
+        if not isinstance(c, str):
+            continue
+        s, e = int(df.iat[i, 1]), int(df.iat[i, 2])
+        n = df.iat[i, 3] if ncol > 3 and df.iat[i, 3] else f"{c}:{s + 1}-{e}"
+        st = STRAND.get(df.iat[i, 5], 0) if ncol > 5 else 0
+        if ncol >= 12:
+            bsz = [int(x) for x in df.iat[i, 10].strip(",").split(",") if x]
+            bst = [int(x) for x in df.iat[i, 11].strip(",").split(",") if x]
+            for a, b in zip(bst, bsz, strict=True):
+                records.append((c, n, st, i, s + a, s + a + b))
+        else:
+            records.append((c, n, st, i, s, e))
+    return collapse(records)
+
+
+def read_gene_models(path: Path, sizes: Mapping[str, int]) -> dict[str, list[GeneModel]]:
+    """Dispatch on the file name: .gtf, .bed, otherwise a genePred table (all optionally .gz)."""
+    name = Path(path).name.lower()
+    if name.endswith(".gz"):
+        name = name[:-3]
+    if name.endswith(".gtf"):
+        return _read_gtf(Path(path), sizes)
+    if name.endswith((".gff", ".gff3")):
+        raise ValueError(f"{path}: GFF3 gene models are not supported; convert to GTF first (for example `gffread -T`)")
+    if name.endswith(".bed"):
+        return _read_bed(Path(path), sizes)
+    return _read_genepred(Path(path), sizes)
+
+
+def _download(url: str, dest: Path) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_name(dest.name + ".part")
+    with urllib.request.urlopen(url, timeout=120) as r, open(tmp, "wb") as fh:
+        shutil.copyfileobj(r, fh)
+    tmp.replace(dest)
+
+
+def refseq_table(genome: str, *, log: Log = print, cache: Path | None = None) -> Path:
+    """The UCSC refGene table for `genome`, downloaded once into the cache."""
+    if genome not in REFSEQ_GENOMES:
+        raise ValueError(f"no RefSeq gene table for genome {genome!r}; set `genes:` to a GTF, genePred or BED file")
+    dest = (cache or cache_dir()) / "genomes" / genome / "refGene.txt.gz"
+    if not dest.exists():
+        url = REFSEQ_URL.format(genome=genome)
+        log(f"  tracks: downloading RefSeq gene models {url} -> {dest}")
+        _download(url, dest)
+        log(f"  tracks: cached {dest.stat().st_size / 1e6:.1f} MB")
+    return dest
+
+
+def gene_models_for(genes: str | None, genome: str, sizes: Mapping[str, int], *, log: Log = print) -> tuple[dict[str, list[GeneModel]], str] | None:
+    """Models and a source label for a `tracks.genes` setting, or None when there is nothing to draw."""
+    if genes is None or genes == "none":
+        return None
+    if genes == "refseq":
+        if genome not in REFSEQ_GENOMES:
+            log(f"  tracks: no RefSeq gene table for genome {genome!r}; set `genes:` to a GTF, genePred or BED for a gene track")
+            return None
+        return read_gene_models(refseq_table(genome, log=log), sizes), f"RefSeq ({genome})"
+    path = Path(genes)
+    return read_gene_models(path, sizes), path.name
+
+
+def write_gene_models(models: Mapping[str, list[GeneModel]], out_dir: Path, *, source: str, order: Iterable[str] | None = None) -> int:
+    """index.json plus one <chrom>.json per chromosome; returns the number of genes written."""
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    chroms = [c for c in (order or models) if c in models] if order else list(models)
+    index = {"source": source, "chroms": {c: [g.name for g in models[c]] for c in chroms}}
+    (out_dir / "index.json").write_text(json.dumps(index, separators=(",", ":")))
+    n = 0
+    for c in chroms:
+        (out_dir / f"{c}.json").write_text(json.dumps([g.row() for g in models[c]], separators=(",", ":")))
+        n += len(models[c])
+    return n

spatialscape/geometry.py

@@ -0,0 +1,132 @@
+"""Coordinate frames: source units -> microns, flips/rotations shared by points and images."""
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import numpy as np
+
+from .config import SampleSpec
+from .readers import SampleInput
+
+VISIUM_SPOT_DIAMETER_UM = 55.0
+VISIUM_SPOT_SPACING_UM = 100.0
+
+
+def _scalar(v: Any) -> float:
+    return float(np.asarray(v).ravel()[0])
+
+
+def _no_spot_rule_for_hd(spec: SampleSpec, sf: dict[str, Any]) -> None:
+    """The 55 um spot / 100 um spacing rule is Visium geometry; applied to Visium HD bins it inflates the scale
+    by the bin pitch / 100 um, silently. HD needs SpaceRanger's microns_per_pixel."""
+    if spec.platform == "visium_hd" and "microns_per_pixel" not in sf:
+        raise ValueError(
+            f"sample {spec.id}: Visium HD cannot use the Visium spot-spacing rule; set `microns: {{ microns_per_unit: <SpaceRanger microns_per_pixel> }}`"
+            " (bin coordinates in full-resolution pixels) or `already_microns: true`"
+        )
+
+
+def _from_scalefactors(sf: dict[str, Any]) -> float | None:
+    if "microns_per_pixel" in sf:
+        return float(sf["microns_per_pixel"])
+    if "spot_diameter_fullres" in sf:
+        return VISIUM_SPOT_DIAMETER_UM / float(sf["spot_diameter_fullres"])
+    return None
+
+
+def microns_per_unit(spec: SampleSpec, sinput: SampleInput) -> float:
+    """Microns per coordinate unit of `sinput.coords`."""
+    m = spec.microns
+    if m.already_microns:
+        return 1.0
+    if m.microns_per_unit is not None:
+        return float(m.microns_per_unit)
+    if m.spot_diameter_fullres is not None:
+        return VISIUM_SPOT_DIAMETER_UM / float(m.spot_diameter_fullres)
+    if m.spot_spacing is not None:
+        return VISIUM_SPOT_SPACING_UM / float(m.spot_spacing)
+    if m.scalefactors_json is not None:
+        with open(m.scalefactors_json) as fh:
+            sf = json.load(fh)
+        _no_spot_rule_for_hd(spec, sf)
+        v = _from_scalefactors(sf)
+        if v is None:
+            raise ValueError(f"{m.scalefactors_json}: no microns_per_pixel / spot_diameter_fullres")
+        return v
+    # auto
+    if spec.kind == "embedding":
+        ext = np.ptp(sinput.coords, axis=0).max()
+        return float(spec.extent / ext) if ext > 0 else 1.0
+    if spec.platform in ("xenium", "merfish"):
+        return 1.0
+    uns = sinput.uns
+    if "spot_nn_spacing_level0_px" in uns:
+        return VISIUM_SPOT_SPACING_UM / _scalar(uns["spot_nn_spacing_level0_px"])
+    spatial = uns.get("spatial")
+    if isinstance(spatial, dict):
+        for lib in spatial.values():
+            if isinstance(lib, dict) and isinstance(lib.get("scalefactors"), dict):
+                _no_spot_rule_for_hd(spec, lib["scalefactors"])
+                v = _from_scalefactors(lib["scalefactors"])
+                if v is not None:
+                    return v
+    raise ValueError(
+        f"sample {spec.id}: cannot infer microns per coordinate unit; set `microns:` "
+        "(already_microns | microns_per_unit | spot_diameter_fullres | spot_spacing | scalefactors_json). "
+        "Coordinates already in microns (Xenium, MERFISH, most segmentation outputs): `platform: xenium` / `merfish`, "
+        "or `microns: { already_microns: true }`; Visium HD: `microns: { microns_per_unit: <SpaceRanger microns_per_pixel> }`"
+    )
+
+
+def default_point_radius(spec: SampleSpec, xy_um: np.ndarray) -> float:
+    if spec.point_radius is not None:
+        return float(spec.point_radius)
+    if spec.kind == "embedding":
+        return float(0.004 * np.ptp(xy_um, axis=0).max())
+    if spec.platform == "visium":
+        return VISIUM_SPOT_DIAMETER_UM / 2
+    if spec.platform == "visium_hd":
+        return max(1.0, 0.5 * _median_nn_distance(xy_um))
+    return 5.0
+
+
+def _median_nn_distance(xy: np.ndarray, sample: int = 2000, seed: int = 0) -> float:
+    from scipy.spatial import cKDTree
+
+    rng = np.random.default_rng(seed)
+    idx = rng.choice(len(xy), size=min(sample, len(xy)), replace=False)
+    tree = cKDTree(xy)
+    d, _ = tree.query(xy[idx], k=2)
+    return float(np.median(d[:, 1]))
+
+
+def transform_points(xy: np.ndarray, width: float, height: float, flip: str, rotate: int) -> np.ndarray:
+    """Apply flip then rotate in a frame of size (width, height). Mirrors transform_image."""
+    x, y = xy[:, 0].copy(), xy[:, 1].copy()
+    w, h = float(width), float(height)
+    if flip == "x":
+        x = w - x
+    elif flip == "y":
+        y = h - y
+    if rotate == 90:  # counter-clockwise, like np.rot90(k=1) on a (y, x) image
+        x, y = y, w - x
+        w, h = h, w
+    elif rotate == 180:
+        x, y = w - x, h - y
+    elif rotate == 270:
+        x, y = h - y, x
+        w, h = h, w
+    return np.stack([x, y], axis=1)
+
+
+def transform_image(img: np.ndarray, flip: str, rotate: int) -> np.ndarray:
+    """img is (c, y, x). Same geometric operation as transform_points."""
+    if flip == "x":
+        img = img[:, :, ::-1]
+    elif flip == "y":
+        img = img[:, ::-1, :]
+    k = {0: 0, 90: 1, 180: 2, 270: 3}[rotate]
+    if k:
+        img = np.rot90(img, k=k, axes=(1, 2))
+    return np.ascontiguousarray(img)

spatialscape/images.py

@@ -0,0 +1,240 @@
+"""Read H&E / fluorescence images and write OME-Zarr (NGFF 0.4, zarr v2) pyramids Viv can read."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import zarr
+from numcodecs import Blosc
+
+from .geometry import transform_image
+
+TILE = 512
+
+
+def read_image(path: Path, *, rgb: bool = True, max_size: int | None = None) -> np.ndarray:
+    """Return (c, y, x). Supports OME-Zarr dirs, tiff, png/jpg. With `max_size`, a pyramidal source (OME-Zarr
+    multiscales, pyramidal/OME-TIFF) is read at the smallest level that is still at least that large, so a 40x
+    scan is never decoded in full just to be shrunk."""
+    path = Path(path)
+    if path.is_dir():
+        return _read_ome_zarr(path, max_size=max_size)
+    suf = path.suffix.lower()
+    if suf in (".tif", ".tiff", ".btf"):
+        import tifffile
+
+        with tifffile.TiffFile(str(path)) as tf:
+            series = tf.series[0]
+            level = _pick_level([tuple(lv.shape) for lv in series.levels], series.axes, max_size) if max_size and len(series.levels) > 1 else 0
+            arr = np.asarray(series.levels[level].asarray()) if level else np.asarray(series.asarray())
+    else:
+        import imageio.v3 as iio
+
+        arr = np.asarray(iio.imread(path))
+    return _to_cyx(arr, rgb=rgb)
+
+
+def _pick_level(shapes: list[tuple[int, ...]], axes: str | list[str], max_size: int | None) -> int:
+    """The last pyramid level whose longer side is still >= max_size (finer levels add nothing after limit_size)."""
+    if not max_size:
+        return 0
+    axes = [a.lower() for a in axes]
+    yi = axes.index("y") if "y" in axes else -2
+    xi = axes.index("x") if "x" in axes else -1
+    pick = 0
+    for i, shape in enumerate(shapes):
+        if max(shape[yi], shape[xi]) >= max_size:
+            pick = i
+        else:
+            break
+    return pick
+
+
+def _to_cyx(arr: np.ndarray, *, rgb: bool = True) -> np.ndarray:
+    """(c, y, x). A 4th plane is alpha only for an RGB image; a 4-channel fluorescence stack keeps all four."""
+    arr = np.squeeze(arr)
+    if arr.ndim == 2:
+        return arr[None, :, :]
+    if arr.ndim == 3:
+        if arr.shape[-1] in (3, 4) and arr.shape[0] not in (3, 4):
+            arr = np.moveaxis(arr, -1, 0)
+        if arr.shape[0] == 4 and rgb:
+            arr = arr[:3]
+        return np.ascontiguousarray(arr)
+    raise ValueError(f"cannot interpret image with shape {arr.shape}")
+
+
+def _read_ome_zarr(path: Path, *, max_size: int | None = None) -> np.ndarray:
+    g = zarr.open_group(str(path), mode="r")
+    attrs = dict(g.attrs)
+    ms = attrs.get("multiscales") or attrs.get("ome", {}).get("multiscales")
+    if not ms:
+        raise ValueError(f"{path}: no multiscales metadata")
+    axes = [a["name"] if isinstance(a, dict) else a for a in ms[0].get("axes", [])]
+    paths = [d["path"] for d in ms[0]["datasets"]]
+    level = _pick_level([tuple(g[p].shape) for p in paths], axes or ["c", "y", "x"][-len(g[paths[0]].shape) :], max_size) if max_size else 0
+    arr = np.asarray(g[paths[level]][:])
+    if axes:
+        # drop t/z if singleton, then order to c,y,x
+        order = []
+        for want in ("c", "y", "x"):
+            if want in axes:
+                order.append(axes.index(want))
+        keep = [i for i, a in enumerate(axes) if a in ("c", "y", "x")]
+        drop = [i for i in range(arr.ndim) if i not in keep]
+        for i in sorted(drop, reverse=True):
+            if arr.shape[i] != 1:
+                raise ValueError(f"{path}: non-singleton axis {axes[i]} unsupported")
+            arr = np.take(arr, 0, axis=i)
+        axes2 = [axes[i] for i in keep]
+        arr = np.transpose(arr, [axes2.index(a) for a in ("c", "y", "x") if a in axes2])
+    return _to_cyx(arr)
+
+
+def downsample2(img: np.ndarray, strip: int = 1024) -> np.ndarray:
+    """(c, y, x) -> 2x mean downsample, keeping dtype. Integer images are summed in uint32 strip by strip, so the
+    working memory is a slice of the output, not a float32 copy of the whole image."""
+    c, h, w = img.shape
+    h2, w2 = h // 2, w // 2
+    if h2 == 0 or w2 == 0:
+        return img
+    out = np.empty((c, h2, w2), dtype=img.dtype)
+    integer = np.issubdtype(img.dtype, np.integer)
+    for y in range(0, h2, strip):
+        y1 = min(h2, y + strip)
+        blk = img[:, 2 * y : 2 * y1, : 2 * w2]
+        if integer:
+            acc = blk[:, 0::2, 0::2].astype(np.uint32)
+            acc += blk[:, 1::2, 0::2]
+            acc += blk[:, 0::2, 1::2]
+            acc += blk[:, 1::2, 1::2]
+            out[:, y:y1] = ((acc + 2) >> 2).astype(img.dtype)
+        else:
+            out[:, y:y1] = blk.astype(np.float32).reshape(c, y1 - y, 2, w2, 2).mean(axis=(2, 4)).astype(img.dtype)
+    return out
+
+
+def limit_size(img: np.ndarray, max_size: int | None) -> np.ndarray:
+    if not max_size:
+        return img
+    while max(img.shape[1], img.shape[2]) > max_size:
+        img = downsample2(img)
+    return img
+
+
+def write_ome_zarr(
+    img: np.ndarray,
+    out_path: Path,
+    *,
+    rgb: bool,
+    channel_names: list[str] | None = None,
+    channel_colors: list[str] | None = None,
+    windows: list[tuple[float, float]] | None = None,
+    min_level_px: int = 256,
+) -> list[tuple[int, int, int]]:
+    """Write NGFF 0.4 (zarr v2) with relative pyramid scales; returns level shapes."""
+    c = img.shape[0]
+    levels = [img]
+    while max(levels[-1].shape[1], levels[-1].shape[2]) > min_level_px:
+        nxt = downsample2(levels[-1])
+        if nxt.shape == levels[-1].shape:
+            break
+        levels.append(nxt)
+    store = zarr.storage.LocalStore(str(out_path))
+    root = zarr.open_group(store, mode="w", zarr_format=2)
+    comp = Blosc(cname="zstd", clevel=5, shuffle=Blosc.SHUFFLE)
+    datasets = []
+    for i, lv in enumerate(levels):
+        arr = root.create_array(
+            str(i),
+            shape=lv.shape,
+            chunks=(1, min(TILE, lv.shape[1]), min(TILE, lv.shape[2])),
+            dtype=lv.dtype,
+            compressors=comp,
+            fill_value=0,
+        )
+        arr[:] = lv
+        datasets.append({"path": str(i), "coordinateTransformations": [{"type": "scale", "scale": [1, 2**i, 2**i]}]})
+    names = channel_names or ([f"{k}" for k in ("R", "G", "B")] if rgb else [f"ch{k}" for k in range(c)])
+    colors = channel_colors or (["ff0000", "00ff00", "0000ff"] if rgb else ["ffffff"] * c)
+    maxv = 255 if img.dtype == np.uint8 else int(np.iinfo(img.dtype).max) if np.issubdtype(img.dtype, np.integer) else 1
+    wins = windows or [(0, maxv)] * c
+    root.attrs.update(
+        {
+            "multiscales": [
+                {
+                    "version": "0.4",
+                    "name": out_path.stem,
+                    "axes": [
+                        {"name": "c", "type": "channel"},
+                        {"name": "y", "type": "space"},
+                        {"name": "x", "type": "space"},
+                    ],
+                    "datasets": datasets,
+                    "type": "mean",
+                }
+            ],
+            "omero": {
+                "id": 1,
+                "name": out_path.stem,
+                "version": "0.4",
+                "rdefs": {"model": "color" if rgb else "greyscale", "defaultT": 0, "defaultZ": 0},
+                "channels": [
+                    {
+                        "label": names[k],
+                        "color": colors[k].lstrip("#"),
+                        "active": True,
+                        "window": {"min": 0, "max": maxv, "start": wins[k][0], "end": wins[k][1]},
+                    }
+                    for k in range(c)
+                ],
+            },
+        }
+    )
+    return [tuple(lv.shape) for lv in levels]
+
+
+def prepare_image(path: Path, flip: str, rotate: int, max_size: int | None, *, rgb: bool = True) -> np.ndarray:
+    img = read_image(path, rgb=rgb, max_size=max_size)
+    img = limit_size(img, max_size)
+    return transform_image(img, flip, rotate)
+
+
+def read_image_shape(path: Path, *, rgb: bool = True) -> tuple[int, int, int]:
+    """(c, y, x) at level 0 from the file header: TIFF series metadata, PNG/JPEG headers, OME-Zarr attributes."""
+    path = Path(path)
+    suffix = path.suffix.lower()
+    if suffix in (".tif", ".tiff"):
+        import tifffile
+
+        with tifffile.TiffFile(str(path)) as tf:
+            series = tf.series[0]
+            shape, axes = tuple(series.shape), str(series.axes)
+        if {"Y", "X"} <= set(axes):
+            c = shape[axes.index("C")] if "C" in axes else shape[axes.index("S")] if "S" in axes else 1
+            return (3 if c == 4 and rgb else c, shape[axes.index("Y")], shape[axes.index("X")])
+        return _to_cyx(np.zeros(shape, dtype=np.uint8), rgb=rgb).shape
+    if suffix in (".png", ".jpg", ".jpeg"):
+        from PIL import Image
+
+        with Image.open(path) as im:
+            w, h = im.size
+            c = len(im.getbands())
+        return (3 if c == 4 and rgb else c, h, w)
+    if path.is_dir():
+        g = zarr.open_group(str(path), mode="r")
+        attrs = dict(g.attrs)
+        ms = attrs.get("multiscales") or attrs.get("ome", {}).get("multiscales")
+        shape = tuple(g[ms[0]["datasets"][0]["path"]].shape)
+        axes = [a["name"] if isinstance(a, dict) else a for a in ms[0].get("axes", [])]
+        if axes and set(axes) >= {"y", "x"}:
+            c = shape[axes.index("c")] if "c" in axes else 1
+            return (c, shape[axes.index("y")], shape[axes.index("x")])
+        return _to_cyx(np.zeros(shape, dtype=np.uint8), rgb=rgb).shape
+    return read_image(path, rgb=rgb).shape
+
+
+def json_dump(obj, path: Path) -> None:
+    with open(path, "w") as fh:
+        json.dump(obj, fh, indent=1)