VisionLanguageGroup/MicroscopyMatching
0
1"""Image standardization.2 3``standardize_image`` collapses any microscopy TIFF variant (16-bit/float,4multi-channel, multi-page Z/time stacks) to one canonical 8-bit RGB PNG, used5for both the preview and the model.6"""7 8import os9import tempfile10from contextlib import contextmanager11from typing import NamedTuple12 13import numpy as np14from PIL import Image15 16try:17 import tifffile18except ImportError: # pragma: no cover - tifffile is a project dependency19 tifffile = None20 21# The model resizes any input to 512x512, so a larger image loses detail rather22# than adding any.23RECOMMENDED_SIZE = 51224WARN_SIZE = 307225MAX_SIZE = 409626# Warn below this: a shorter side under 32px means >16x upscaling to 512, so27# there is almost no real detail for the model to work with.28MIN_SIZE = 3229# An integer image is a low-dynamic-range candidate when its full span is tiny30# relative to the dtype's max. Float has no fixed range, so it is skipped.31NARROW_RANGE_FRAC = 0.0232# Sparse fluorescence can have a narrow span but still be valid; only flag it33# when the max is not meaningfully above the 99th percentile.34BRIGHT_TAIL_FRAC = 0.135# Measured on the uncompressed array, not the file size on disk: compression36# makes disk size a poor proxy for what opening the file actually allocates.37MAX_READ_BYTES = 300 * 1024 ** 238 39# Axis roles, per tifffile's `series.axes` naming. Anything else (Z, T, ...) is a40# stack axis, indexed by frame.41_CHANNEL = ("C", "S") # C = separate channel planes, S = interleaved RGB samples42_UNKNOWN = ("Q", "I") # file named no axis; only these may be *guessed* as channels43 44 45TIFF_EXTENSIONS = (46 sorted(47 {".ome.tif", ".ome.tiff"}48 | {"." + e for e in tifffile.TIFF.FILE_EXTENSIONS if "." not in e}49 )50 if tifffile is not None else [".ome.tif", ".ome.tiff", ".tif", ".tiff"]51)52 53 54def _series_planes(series):55 """Number of 2D planes in a series: the product of every axis that is not56 spatial (Y/X) or RGB samples (S). Axes-aware on purpose - a positional57 shape[:-2] would miscount RGB, which tifffile stores as a trailing S axis."""58 planes = int(np.prod([59 dim for axis, dim in zip(str(series.axes), series.shape)60 if axis not in ("Y", "X", "S")61 ]))62 return planes or 163 64 65def _ome_needs_single_file(tif):66 """True for a multi-file OME-TIFF whose logical series spans more planes than67 this file holds: its data lives in sibling files that are not part of a68 single upload, so the assembled series would zero-fill the out-of-file69 planes. A self-contained OME has all its planes in-file (planes == pages)."""70 try:71 return bool(tif.is_ome) and _series_planes(tif.series[0]) > len(tif.pages)72 except Exception: # noqa: BLE001 - detection must never break the read73 return False74 75 76@contextmanager77def _open_tiff(path):78 """Open a TIFF; for a multi-file OME set whose series spans beyond this file,79 reopen just this file's own pages (is_ome=False) so out-of-file planes are80 not zero-filled. Detection is lazy (shape/pages only, no pixel decode)."""81 tif = tifffile.TiffFile(path)82 try:83 if _ome_needs_single_file(tif):84 tif.close()85 tif = tifffile.TiffFile(path, is_ome=False)86 yield tif87 finally:88 tif.close()89 90 91def _read_tiff(path):92 """Read a TIFF-family file into (array, axes). Raises if not readable."""93 with _open_tiff(path) as tif:94 series = tif.series[0]95 page = tif.pages[0]96 arr = np.asarray(series.asarray())97 axes = str(series.axes)98 99 is_palette = (getattr(page, "photometric", None) == tifffile.PHOTOMETRIC.PALETTE100 and not any(a in _CHANNEL for a in axes))101 if is_palette and page.colormap is not None:102 colormap = np.asarray(page.colormap) # (3, 2**bits), uint16103 rgb = np.moveaxis(colormap[:, arr], 0, -1) # (..., H, W, 3)104 # TIFF colormaps are 16-bit; scale down to 8-bit (65535/255=257).105 arr = np.round(rgb / 257.0).astype(np.uint8)106 axes = axes + "S" # the LUT added an RGB sample axis107 return arr, axes108 109 110def _read_array(path):111 """Load an image file into (array, axes) where axes names each dimension.112 113 axes uses tifffile's convention ('C' channel, 'S' RGB samples, 'Z'/'T'114 stack, 'Y'/'X' spatial, 'Q' unknown). Indexed / palette images (ImageJ115 "8-bit Color", palette PNG/GIF) are expanded through their color lookup116 table so we return true RGB, not bare indices.117 """118 if tifffile is not None:119 try:120 return _read_tiff(path)121 except Exception: # noqa: BLE001 - not a TIFF, or tifffile cannot parse it122 pass123 124 img = Image.open(path)125 if img.mode in ("P", "PA"):126 img = img.convert("RGB")127 arr = np.asarray(img)128 return arr, ("YXS" if arr.ndim == 3 else "YX")129 130 131def _shape_axes(path):132 """Return (shape, axes) from metadata only - no pixel decode."""133 if tifffile is not None:134 try:135 with _open_tiff(path) as tif:136 series = tif.series[0]137 return tuple(series.shape), str(series.axes)138 except Exception: # noqa: BLE001 - fall through to PIL139 pass140 with Image.open(path) as img:141 w, h = img.size142 if img.mode in ("P", "PA", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV"):143 return (h, w, len(img.getbands())), "YXS"144 return (h, w), "YX"145 146 147_PIL_BITS = {"1": 1, "L": 8, "P": 8, "LA": 8, "PA": 8, "RGB": 8, "RGBA": 8,148 "CMYK": 8, "YCbCr": 8, "LAB": 8, "HSV": 8,149 "I;16": 16, "I;16B": 16, "I;16L": 16, "I": 32, "F": 32}150 151 152def bit_depth(path):153 """Bits per sample, read from metadata only (0 if it cannot be determined)."""154 if tifffile is not None:155 try:156 with _open_tiff(path) as tif:157 return int(np.dtype(tif.series[0].dtype).itemsize) * 8158 except Exception: # noqa: BLE001 - fall through to PIL159 pass160 try:161 with Image.open(path) as img:162 return int(_PIL_BITS.get(img.mode, 0))163 except Exception: # noqa: BLE001164 return 0165 166 167def _plan_axes(shape, axes):168 """Drop size-1 axes, infer a channel axis when the file names none, and move169 it last.170 171 Returns (shape, axes, guessed) as lists + a flag saying whether the channel172 axis had to be inferred rather than read from the file.173 """174 axes = list(axes)175 shape = list(shape)176 if len(axes) != len(shape): # be defensive; keep the trailing spatial axes177 axes = ["Q"] * (len(shape) - 2) + ["Y", "X"]178 179 pairs = [(a, d) for a, d in zip(axes, shape) if d != 1]180 axes = [a for a, _ in pairs]181 shape = [d for _, d in pairs]182 183 guessed = False184 if not any(a in _CHANNEL for a in axes):185 # Only guess on axes the file left unnamed. An axis the file explicitly186 # calls Z/T is a stack even when its length happens to be 3.187 for i, (a, d) in enumerate(zip(axes, shape)):188 if a in _UNKNOWN and d in (2, 3, 4):189 axes[i] = "C"190 guessed = True191 break192 193 ci = next((i for i, a in enumerate(axes) if a in _CHANNEL), None)194 if ci is not None:195 axes.append(axes.pop(ci))196 shape.append(shape.pop(ci))197 return shape, axes, guessed198 199 200def _reduce_to_hwc(arr, axes, frame=0, sub_frame=None):201 """Collapse a labelled array to 2D (H, W) or 3D (H, W, C).202 203 The channel axis (named by the file, or inferred by ``_plan_axes`` only when204 the file names none) is moved last and kept whole. Stack axes (Z / time /205 page) are handled in order:206 * the first is indexed by ``frame`` (the frame slider);207 * the second is indexed by ``sub_frame`` if given, else collapsed by a208 maximum-intensity projection (the natural default for a focal stack -209 keeps the brightest signal across planes rather than an arbitrary one);210 * any deeper ones are projected too.211 """212 axes = list(axes)213 214 for i in range(len(axes) - 1, -1, -1):215 if arr.shape[i] == 1:216 arr = arr.reshape(arr.shape[:i] + arr.shape[i + 1:])217 axes.pop(i)218 219 _, planned, _ = _plan_axes(arr.shape, axes)220 221 if "C" in planned and not any(a in _CHANNEL for a in axes):222 for i, (a, d) in enumerate(zip(axes, arr.shape)):223 if a in _UNKNOWN and d in (2, 3, 4):224 axes[i] = "C"225 break226 ci = next((i for i, a in enumerate(axes) if a in _CHANNEL), None)227 if ci is not None:228 arr = np.moveaxis(arr, ci, -1)229 axes.append(axes.pop(ci))230 231 target = 3 if ci is not None else 2232 stack = 0233 while len(axes) > target:234 if stack == 0:235 idx = max(0, min(int(frame), arr.shape[0] - 1))236 arr = arr[idx]237 elif stack == 1 and sub_frame is not None:238 idx = max(0, min(int(sub_frame), arr.shape[0] - 1))239 arr = arr[idx]240 else:241 arr = arr.max(axis=0) # maximum-intensity projection over this axis242 axes.pop(0)243 stack += 1244 245 return arr246 247 248# tifffile axis letters -> words a microscopist uses, for messages and slider249# labels. Anything unmapped (an unnamed 'Q'/'I' axis, a raw page axis) is just a250# "frame".251_AXIS_LABELS = {"T": "timepoint", "Z": "z-plane"}252 253 254def _axis_label(a):255 return _AXIS_LABELS.get(a, "frame")256 257 258class ImageInfo(NamedTuple):259 """How a file's dimensions were interpreted (read from metadata only).260 261 frames - length of the first stack (Z/T) axis, 1 if not a stack262 channels - length of the channel axis, 1 if single-channel263 axes - the file's own axes string, e.g. 'CZYX' ('Q' = unnamed)264 guessed - True if the channel axis was inferred rather than read265 shape - the file's raw shape as stored, e.g. (1, 4, 1, 1024, 1024)266 width - pixels along the X axis (0 if unknown)267 height - pixels along the Y axis (0 if unknown)268 sub_frames- length of a *second* stack axis (e.g. Z in a time+Z file), 1 if269 there is only one stack axis270 frame_label / sub_label - the words for those two axes ('timepoint', ...)271 """272 frames: int273 channels: int274 axes: str275 guessed: bool276 shape: tuple277 width: int278 height: int279 sub_frames: int = 1280 frame_label: str = "frame"281 sub_label: str = "frame"282 283 284def inspect_image(path):285 """Describe a file's structure from metadata only (no pixel decode)."""286 try:287 shape, axes = _shape_axes(path)288 planned_shape, planned_axes, guessed = _plan_axes(shape, axes)289 has_c = bool(planned_axes) and planned_axes[-1] in _CHANNEL290 channels = int(planned_shape[-1]) if has_c else 1291 292 stack = [(a, d) for a, d in zip(planned_axes, planned_shape)293 if a not in ("Y", "X") and a not in _CHANNEL]294 frames = int(stack[0][1]) if stack else 1295 frame_label = _axis_label(stack[0][0]) if stack else "frame"296 sub_frames = int(stack[1][1]) if len(stack) > 1 else 1297 sub_label = _axis_label(stack[1][0]) if len(stack) > 1 else "frame"298 299 # Spatial size comes from the named Y/X axes, so it stays correct300 # whatever order the other axes are in.301 width = height = 0302 for a, d in zip(axes, shape):303 if a == "Y":304 height = int(d)305 elif a == "X":306 width = int(d)307 if not (width and height): # no Y/X named; fall back to the header probe308 width, height = image_size(path)309 310 return ImageInfo(frames, channels, str(axes), guessed, tuple(shape),311 width, height, sub_frames, frame_label, sub_label)312 except Exception: # noqa: BLE001313 return ImageInfo(1, 1, "", False, (), 0, 0)314 315 316class PixelReport(NamedTuple):317 """Pixel-level sanity checks on decoded image data (see ``pixel_stats``).318 319 decoded - False if the pixels could not be read (corrupt / truncated file)320 finite - False if the frame contains NaN/inf (only checked for float data)321 vmin/vmax - min and max pixel value of the frame322 dtype_max - np.iinfo(dtype).max for integer data, else 0.0323 low_range - True if an integer frame's robust span is a tiny fraction of the324 dtype range (low contrast / narrow dynamic range)325 """326 decoded: bool327 finite: bool328 vmin: float329 vmax: float330 dtype_max: float331 low_range: bool332 333 334def pixel_stats(path, frame=0, sub_frame=None):335 """Decode image pixels and report problems (blank / non-finite / corrupt /336 low dynamic range).337 338 Works on the raw reduced frame, not the RGB-padded one, so a 2-channel339 image's zero-padded third channel does not skew the min/max/finite stats.340 Only call once the header size guard has passed, so decoding is bounded.341 """342 try:343 raw, axes = _read_array(path)344 arr = _reduce_to_hwc(raw, axes, frame=frame, sub_frame=sub_frame)345 except Exception: # noqa: BLE001 - unreadable pixels346 return PixelReport(False, True, 0.0, 0.0, 0.0, False)347 348 if np.issubdtype(arr.dtype, np.floating) and not np.isfinite(arr).all():349 return PixelReport(True, False, 0.0, 0.0, 0.0, False)350 351 vmin, vmax = float(arr.min()), float(arr.max())352 is_int = np.issubdtype(arr.dtype, np.integer)353 dtype_max = float(np.iinfo(arr.dtype).max) if is_int else 0.0354 low_range = False355 if is_int and vmax > vmin and dtype_max > 0:356 narrow = (vmax - vmin) < NARROW_RANGE_FRAC * dtype_max357 if narrow:358 # Among narrow images, spare fluorescence (dark background + sparse359 # bright objects) by checking for a bright tail: its max sits far360 # above the 99th percentile. A dense narrow band (values packed in a361 # thin range) has none, so only that is flagged as low dynamic range.362 p99 = float(np.percentile(arr, 99))363 bright_tail = (vmax - p99) / (vmax - vmin)364 low_range = bright_tail < BRIGHT_TAIL_FRAC365 return PixelReport(True, True, vmin, vmax, dtype_max, low_range)366 367 368 369def array_nbytes(path):370 """Bytes that opening this file's full array would allocate.371 372 Computed from shape + dtype in the header - no pixel decode - so it is safe373 to call on a file that is too large to open. Returns 0 if unknown.374 """375 if tifffile is not None:376 try:377 with _open_tiff(path) as tif:378 series = tif.series[0]379 return int(np.prod(series.shape)) * int(np.dtype(series.dtype).itemsize)380 except Exception: # noqa: BLE001 - fall through to PIL381 pass382 try:383 with Image.open(path) as img:384 w, h = img.size385 return int(w) * int(h) * len(img.getbands()) # 8-bit assumption386 except Exception: # noqa: BLE001387 return 0388 389 390def image_size(path):391 """Return an image's (width, height) by reading only its header.392 393 Never decodes pixel data, so this is safe to call as a size guard on a file394 that would be too large to load. Returns (0, 0) if the size cannot be395 determined, so callers treat it as "unknown" and proceed.396 """397 if tifffile is not None:398 try:399 with tifffile.TiffFile(path) as tif:400 page = tif.pages[0]401 return int(page.imagewidth), int(page.imagelength)402 except Exception: # noqa: BLE001 - fall through to PIL403 pass404 try:405 with Image.open(path) as img: # PIL parses the header lazily406 return int(img.size[0]), int(img.size[1])407 except Exception: # noqa: BLE001408 return 0, 0409 410 411def _to_rgb(arr):412 """Turn a 2D or (H, W, C) array into exactly 3 channels."""413 if arr.ndim == 2:414 return np.stack([arr] * 3, axis=-1)415 416 channels = arr.shape[2]417 if channels == 1:418 return np.repeat(arr, 3, axis=2)419 if channels == 2:420 # Pad a zero third channel rather than inventing signal.421 return np.concatenate([arr, np.zeros_like(arr[:, :, :1])], axis=2)422 return arr[:, :, :3]423 424 425def _load_rgb(path, frame=0, sub_frame=None):426 """Read a file and reduce it to an (H, W, 3) array plus its source dtype."""427 raw, axes = _read_array(path)428 arr = _reduce_to_hwc(raw, axes, frame=frame, sub_frame=sub_frame)429 arr = _to_rgb(arr)430 return arr, raw.dtype431 432 433def _to_uint8(arr, stretch):434 """Map pixel values to uint8.435 436 When ``stretch`` is True (non-8-bit input) a 1st-99th percentile auto-437 contrast stretch is applied - outlier-robust, so a hot/saturated pixel does438 not collapse the visible signal to black. Otherwise values are only clipped,439 so standard 8-bit images are unchanged.440 """441 arr = arr.astype(np.float32)442 443 if not stretch:444 return np.clip(arr, 0, 255).astype(np.uint8)445 446 lo, hi = np.percentile(arr, (1, 99))447 if hi <= lo: # near-flat image; fall back to full min/max448 lo, hi = float(arr.min()), float(arr.max())449 if hi <= lo: # truly constant image450 return np.zeros(arr.shape, dtype=np.uint8)451 452 arr = np.clip((arr - lo) / (hi - lo), 0.0, 1.0)453 return (arr * 255.0).astype(np.uint8)454 455 456def _save_png(arr, out_path=None, out_dir=None, base=None, suffix="_std.png"):457 """Save an (H, W, 3) uint8 array as a PNG and return the path."""458 if out_path is None:459 if out_dir is not None:460 os.makedirs(out_dir, exist_ok=True)461 out_path = os.path.join(out_dir, (base or "image") + suffix)462 else:463 tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)464 out_path = tmp.name465 tmp.close()466 Image.fromarray(arr, mode="RGB").save(out_path)467 return out_path468 469 470def standardize_image(path, frame=0, sub_frame=None, out_dir=None, out_path=None):471 """Standardize any image/TIFF to a canonical 8-bit RGB PNG.472 473 Used for both the preview and the model: one frame, <=3 channels, with a474 1st-99th percentile auto-contrast stretch for 16-bit/float input (8-bit is475 passed through unchanged).476 477 Args:478 path: path to the input image (TIFF, PNG, JPG, ...).479 frame: which frame of the first stack axis to use (0-based; ignored for480 non-stacks).481 sub_frame: which plane of a second stack axis (e.g. Z in a time+Z file)482 to use (0-based). None means combine those planes by a maximum-483 intensity projection.484 out_dir: optional directory for the output PNG.485 out_path: optional explicit output file path (overrides out_dir).486 487 Returns:488 Path to the standardized PNG. On any failure the original ``path`` is489 returned unchanged so callers degrade gracefully.490 """491 try:492 arr, dtype = _load_rgb(path, frame=frame, sub_frame=sub_frame)493 arr = _to_uint8(arr, stretch=(dtype != np.uint8))494 base = os.path.splitext(os.path.basename(path))[0]495 return _save_png(arr, out_path=out_path, out_dir=out_dir, base=base, suffix="_std.png")496 except Exception as e: # noqa: BLE001 - never let preprocessing crash a run497 print(f"⚠️ standardize_image failed for {path}: {e}; using original file")498 return path499