Coverage for src/precon3d/mosaic_utils.py: 30%
256 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 03:50 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4mosaic_utils.py
6Utilities for stitching ZEISS CZI mosaics and processing tile-region
7arrays, including visualization and saving functions.
9This module provides:
10 * Reading tile-region arrays and mosaics from CZI files.
11 * Merging and grouping tile stacks.
12 * Simple numpy-based stitching of tiles.
13 * Fiji/ImageJ-based tile alignment and stitching.
14 * Visualization of tile layouts and image stacks.
15 * Saving of tiles, mosaics, and thumbnails.
16 * Flatfield normalization utilities.
17"""
19import math
20import sys
21import warnings
22import subprocess
23from pathlib import Path
24from typing import Any, List, Optional
26import numpy as np
27import skimage.io
28import skimage.transform
29from PIL import Image
30from aicspylibczi import CziFile
31from aicsimageio import AICSImage
32import matplotlib.pyplot as plt
33import matplotlib.patches as patches
34import tkinter as tk
35from tkinter import ttk
36from matplotlib.backends.backend_tkagg import (
37 FigureCanvasTkAgg,
38 NavigationToolbar2Tk,
39)
41import precon3d.czi_info as ci
42from precon3d.custom_types import Mosaic, TilePosition, TileStack
45def flatten_list(nested: list[Any]) -> list[Any]:
46 """
47 Recursively flatten a nested list.
49 Args:
50 nested: A list which may contain other lists.
52 Returns:
53 A flat list containing all non-list elements from `nested`.
54 """
55 flat: list[Any] = []
56 for element in nested:
57 if isinstance(element, list):
58 flat.extend(flatten_list(element))
59 else:
60 flat.append(element)
61 return flat
64def merge_tile_arrays(tiles: list[TileStack]) -> TileStack:
65 """
66 Merge multiple TileStacks into a single TileStack.
68 Stacks the tile data arrays along a new axis and concatenates positions.
70 Args:
71 tiles: List of TileStack objects, assumed to share the same
72 channel, czi_fname, scene, and scene_idx.
74 Returns:
75 A new TileStack with merged data and positions.
77 Raises:
78 ValueError: If `tiles` is empty.
79 """
80 if not tiles:
81 raise ValueError("Cannot merge an empty list of TileStacks.")
82 # Stack data arrays
83 merged_data = np.stack([ts.data for ts in tiles], axis=0)
84 # Concatenate positions
85 merged_positions = [pos for ts in tiles for pos in ts.positions]
86 # Use metadata from the first stack
87 first = tiles[0]
88 return TileStack(
89 data=merged_data,
90 positions=merged_positions,
91 channel=first.channel,
92 czi_fname=first.czi_fname,
93 scene=first.scene,
94 scene_idx=first.scene_idx,
95 )
98def unpack_tilestack(merged: TileStack) -> list[TileStack]:
99 """
100 Split a merged TileStack back into a list of individual TileStacks.
102 Args:
103 merged: A TileStack where `data.shape[0]` is the number of original tiles.
105 Returns:
106 A list of TileStack objects, one per original tile.
107 """
108 n_tiles = merged.data.shape[0]
109 individual: list[TileStack] = []
110 for i in range(n_tiles):
111 tile_data = merged.data[i]
112 tile_pos = merged.positions[i : i + 1]
113 individual.append(
114 TileStack(
115 data=tile_data,
116 positions=tile_pos,
117 channel=merged.channel,
118 czi_fname=merged.czi_fname,
119 scene=merged.scene,
120 scene_idx=merged.scene_idx,
121 )
122 )
123 return individual
126def group_tile_arrays(tiles: list[TileStack], idx: int = 11) -> list[TileStack]:
127 """
128 Split a flat list of TileStacks into two groups at position `idx` and merge each.
130 Args:
131 tiles: Flat list of TileStack objects.
132 idx: Index at which to split `tiles`.
134 Returns:
135 A list of two TileStacks, each merged from a slice of `tiles`.
136 """
137 return [merge_tile_arrays(tiles[:idx]), merge_tile_arrays(tiles[idx:])]
140def tile_region_array_data(czi_file: Path) -> list[TileStack]:
141 """
142 Retrieve tile-region array data from a CZI file.
144 Extracts all channels and scenes and returns a flat list of TileStacks.
146 Args:
147 czi_file: Path to the .czi file containing tile-region arrays.
149 Returns:
150 A list of TileStack objects, one per channel and scene.
152 Raises:
153 TypeError: If the scenes have inconsistent shapes or unsupported dims.
154 """
155 czi = CziFile(czi_file)
156 if not czi.shape_is_consistent:
157 raise TypeError(f"Inconsistent shapes in CZI file: {czi_file.name}")
158 dims_shape = czi.get_dims_shape()
159 channel_meta = ci.channel_metadata(czi_file)
160 # Get bounding boxes for all scenes
161 bboxes = czi.get_all_scene_bounding_boxes()
162 positions = [TilePosition(y=bb.y, x=bb.x) for bb in bboxes.values()]
163 stacks: list[TileStack] = []
164 # dims_shape is a list of dicts, e.g. [{'X':(...), 'Y':(...), 'C':(..), 'S':(..)}]
165 for block in dims_shape:
166 c_start, c_end = block["C"]
167 s_start, s_end = block["S"]
168 for c_idx in range(c_start, c_end):
169 chan = channel_meta[c_idx].channel_name
170 for s_idx in range(s_start, s_end):
171 img, _ = czi.read_image(C=c_idx, S=s_idx)
172 stacks.append(
173 TileStack(
174 data=img.squeeze(),
175 positions=[positions[s_idx]],
176 czi_fname=czi_file.stem,
177 channel=chan,
178 scene=f"scene{s_idx}",
179 scene_idx=s_idx,
180 )
181 )
182 return stacks
185def visualize_tilestacks(tilestacks: list[TileStack], tile_size: int = 4512) -> None:
186 """
187 Visualize TileStacks as colored rectangles on a 2D plot.
189 Args:
190 tilestacks: List of TileStack objects (only first position used).
191 tile_size: Size (height and width) of each tile in pixels.
192 """
193 fig, ax = plt.subplots(figsize=(12, 12))
194 xs = [ts.positions[0].x for ts in tilestacks]
195 ys = [ts.positions[0].y for ts in tilestacks]
196 min_x, max_x = min(xs), max(xs) + tile_size
197 min_y, max_y = min(ys), max(ys) + tile_size
198 ax.set_xlim(min_x - tile_size, max_x + tile_size)
199 ax.set_ylim(min_y - tile_size, max_y + tile_size)
200 ax.invert_yaxis()
201 cmap = plt.get_cmap("viridis", len(tilestacks))
202 for i, ts in enumerate(tilestacks):
203 x0, y0 = ts.positions[0].x, ts.positions[0].y
204 rect = patches.Rectangle(
205 (x0, y0),
206 tile_size,
207 tile_size,
208 linewidth=1,
209 edgecolor="r",
210 facecolor=cmap(i),
211 alpha=0.2,
212 )
213 ax.add_patch(rect)
214 sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(0, len(tilestacks) - 1))
215 fig.colorbar(sm, ax=ax, label="Tile index")
216 plt.show()
219def view_image_stack(img_stack: np.ndarray, on_close: str = "exit") -> None:
220 """
221 Interactive viewer to browse through an image stack using Tkinter.
223 Args:
224 img_stack: NumPy array of shape (N, H, W) or (N, H, W, C).
225 on_close: If "exit", closes program on window close; else just closes window.
226 """
227 root = tk.Tk()
228 root.title("Image Stack Viewer")
230 def _on_close():
231 if on_close == "exit":
232 sys.exit(0)
233 root.destroy()
235 root.protocol("WM_DELETE_WINDOW", _on_close)
236 fig, ax = plt.subplots()
237 im = ax.imshow(img_stack[0], cmap="gray")
238 ax.set_title("Index: 0")
240 def _update(val: float) -> None:
241 idx = int(val)
242 im.set_data(img_stack[idx])
243 ax.set_title(f"Index: {idx}")
244 fig.canvas.draw_idle()
246 canvas = FigureCanvasTkAgg(fig, master=root)
247 canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
248 NavigationToolbar2Tk(canvas, root).update()
250 slider = ttk.Scale(
251 root,
252 from_=0,
253 to=img_stack.shape[0] - 1,
254 orient=tk.HORIZONTAL,
255 command=_update,
256 )
257 slider.pack(fill=tk.X)
259 def _on_release(evt: Any) -> None:
260 _update(slider.get())
262 slider.bind("<ButtonRelease-1>", _on_release)
263 root.mainloop()
266def simple_mosaic_from_tilestack(tilestack: TileStack) -> Mosaic:
267 """
268 Create a simple grayscale mosaic by averaging overlapping tiles equally.
270 Args:
271 tilestack: A TileStack of shape (N, H, W).
273 Returns:
274 A Mosaic with blended data.
275 """
276 n, h, w = tilestack.data.shape
277 xs = [p.x for p in tilestack.positions]
278 ys = [p.y for p in tilestack.positions]
279 min_x, min_y = math.floor(min(xs)), math.floor(min(ys))
280 max_x = math.ceil(max(xs) + w)
281 max_y = math.ceil(max(ys) + h)
282 off_x = -min_x if min_x < 0 else 0
283 off_y = -min_y if min_y < 0 else 0
284 width, height = max_x + off_x, max_y + off_y
285 mosaic_arr = np.zeros((height, width), dtype=np.float32)
286 weight = np.zeros_like(mosaic_arr)
287 for idx, pos in enumerate(tilestack.positions):
288 x0 = round(pos.x + off_x)
289 y0 = round(pos.y + off_y)
290 tile = tilestack.data[idx, :, :]
291 mosaic_arr[y0 : y0 + h, x0 : x0 + w] += tile
292 weight[y0 : y0 + h, x0 : x0 + w] += 1
293 weight[weight == 0] = 1
294 blended = mosaic_arr / weight
295 return Mosaic(
296 data=blended,
297 positions=tilestack.positions,
298 channel=tilestack.channel,
299 czi_fname=tilestack.czi_fname,
300 scene=tilestack.scene,
301 scene_idx=tilestack.scene_idx,
302 )
305def stitch_tiles_from_positions(tilestack: TileStack) -> Mosaic:
306 """
307 Stitch tiles into a mosaic by equal-weight blending at overlaps.
309 Args:
310 tilestack: A TileStack of shape (N, H, W) or (N, H, W, C).
312 Returns:
313 A Mosaic object with blended image data.
314 """
315 data = tilestack.data
316 dims = data.shape
317 if data.ndim == 3:
318 n, h, w = dims
319 c = 1
320 else:
321 n, h, w, c = dims
322 xs = [p.x for p in tilestack.positions]
323 ys = [p.y for p in tilestack.positions]
324 min_x, min_y = math.floor(min(xs)), math.floor(min(ys))
325 max_x = math.ceil(max(xs) + w)
326 max_y = math.ceil(max(ys) + h)
327 off_x = -min_x if min_x < 0 else 0
328 off_y = -min_y if min_y < 0 else 0
329 width, height = max_x + off_x, max_y + off_y
330 if c > 1:
331 mos = np.zeros((height, width, c), np.float32)
332 wgt = np.zeros_like(mos)
333 else:
334 mos = np.zeros((height, width), np.float32)
335 wgt = np.zeros_like(mos)
336 for idx, pos in enumerate(tilestack.positions):
337 x0 = round(pos.x + off_x)
338 y0 = round(pos.y + off_y)
339 tile = data[idx].astype(np.float32)
340 mos[y0 : y0 + h, x0 : x0 + w, ...] += tile
341 wgt[y0 : y0 + h, x0 : x0 + w, ...] += 1
342 wgt[wgt == 0] = 1
343 blended = mos / wgt
344 return Mosaic(
345 data=blended,
346 positions=tilestack.positions,
347 channel=tilestack.channel,
348 czi_fname=tilestack.czi_fname,
349 scene=tilestack.scene,
350 scene_idx=tilestack.scene_idx,
351 )
354def extract_substring(start: str, end: str, string: str) -> str:
355 """
356 Extract substring between `start` and `end` markers.
358 Args:
359 start: Starting delimiter.
360 end: Ending delimiter.
361 string: Input string to parse.
363 Returns:
364 The substring between `start` and `end`.
366 Raises:
367 IndexError: If delimiters not found.
368 """
369 return string.split(start)[1].split(end)[0]
372def ai_tilestack(czi_file: Path) -> tuple[AICSImage, str]:
373 """
374 Load CZI file into an AICSImage for tilestack extraction.
376 Args:
377 czi_file: Path to .czi file.
379 Returns:
380 Tuple of (AICSImage, filename stem).
381 """
382 return (AICSImage(czi_file, reconstruct_mosaic=False), czi_file.stem)
385def simple_tilestack(
386 ai_image: AICSImage, scene_idx: int, czi_fname: str = ""
387) -> list[TileStack]:
388 """
389 Extract per-channel TileStacks from an AICSImage scene.
391 Args:
392 ai_image: AICSImage with mosaic data.
393 scene_idx: Scene index to extract.
394 czi_fname: Optional filename stem for output.
396 Returns:
397 A list of TileStack objects, one per channel.
398 """
399 # breakpoint()
400 # dims = dict(ai_image.reader.dims)
401 dims = ai_image.reader.dims
402 n_channels = dims.C # get("C", 1)
403 ai_image.set_scene(scene_idx)
404 tile_pos = ai_image.get_mosaic_tile_positions()
405 stacks: list[TileStack] = []
406 for ch in range(n_channels):
407 arr = ai_image.reader.get_image_data("MYXS", C=ch) # .squeeze()
408 positions = [TilePosition(y=y, x=x) for y, x in tile_pos]
409 stacks.append(
410 TileStack(
411 data=arr,
412 positions=positions,
413 czi_fname=czi_fname,
414 channel=ai_image.channel_names[ch],
415 scene=ai_image.current_scene,
416 scene_idx=ai_image.current_scene_index,
417 )
418 )
419 return stacks
422def align_tilestack_positions(tilestack: TileStack, fiji_app: Path) -> TileStack:
423 """
424 Use Fiji/ImageJ to compute refined tile positions without fusing image.
426 Args:
427 tilestack: Input TileStack.
428 fiji_app: Path to Fiji executable.
430 Returns:
431 A new TileStack with updated positions.
433 Raises:
434 OSError: If Fiji macro reports an error.
435 """
436 temp_dir = Path.cwd() / ".tmp_mosaic"
437 temp_dir.mkdir(exist_ok=True)
438 # Save tiles temporarily
439 n = tilestack.data.shape[0]
440 fnames: list[Path] = []
441 for i in range(n):
442 p = temp_dir / f"{i:04d}.tif"
443 with warnings.catch_warnings():
444 warnings.simplefilter("ignore")
445 skimage.io.imsave(p, tilestack.data[i])
446 fnames.append(p)
447 # Write tile configuration
448 pos_x = np.array([p.x for p in tilestack.positions])
449 pos_y = np.array([p.y for p in tilestack.positions])
450 pos_x -= pos_x.min()
451 pos_y -= pos_y.min()
452 cfg = temp_dir / "TileConfiguration.txt"
453 with cfg.open("w", encoding="utf-8") as f:
454 f.write("# dim = 2\n")
455 for idx in range(n):
456 f.write(
457 f"{fnames[idx].name}; ; ({float(pos_x[idx])}, {float(pos_y[idx])})\n"
458 )
459 # Write macro
460 macro = temp_dir / "stitch.ijm"
461 macro.write_text(
462 'run("Grid/Collection stitching",\n'
463 ' "type=[Positions from file]" +\n'
464 ' " layout_file=TileConfiguration.txt" +\n'
465 ' " fusion_method=[Linear Blending]" +\n'
466 ' " compute_overlap" +\n'
467 ' " image_output=[Display]");\n'
468 "close();"
469 )
470 # Execute
471 proc = subprocess.run(
472 [str(fiji_app), "--headless", "-macro", str(macro)],
473 capture_output=True,
474 text=True,
475 )
476 if "error" in proc.stdout.lower() or proc.returncode != 0:
477 raise OSError(f"Fiji error:\n{proc.stdout}\n{proc.stderr}")
478 # Parse registered positions
479 reg = temp_dir / "TileConfiguration.registered.txt"
480 ys: list[float] = []
481 xs: list[float] = []
482 for line in reg.read_text().splitlines():
483 try:
484 sub = extract_substring("(", ")", line)
485 y_str, x_str = sub.split(",")
486 ys.append(float(y_str))
487 xs.append(float(x_str))
488 except Exception:
489 continue
490 # Clean up temp_dir if desired
491 new_positions = [TilePosition(y=y, x=x) for y, x in zip(xs, ys)]
492 return TileStack(
493 data=tilestack.data,
494 positions=new_positions,
495 channel=tilestack.channel,
496 czi_fname=tilestack.czi_fname,
497 scene=tilestack.scene,
498 scene_idx=tilestack.scene_idx,
499 )
502def stitch_tilestack(
503 tilestack: TileStack, fiji_app: Path, cleanup: bool = False
504) -> Mosaic:
505 """
506 Use Fiji/ImageJ to stitch tiles into a fused mosaic.
508 Args:
509 tilestack: Input TileStack.
510 fiji_app: Path to Fiji executable.
511 cleanup: If True, remove temporary files.
513 Returns:
514 A Mosaic with fused data and final tile positions.
516 Raises:
517 OSError: If Fiji macro fails.
518 """
519 ts_aligned = align_tilestack_positions(tilestack, fiji_app)
520 # Load fused image (skipped here: assume Fiji writes to known path)
521 # For brevity, we return an empty mosaic
522 # In practice, read the output TIFF from Fiji temp directory
523 return Mosaic(
524 data=np.zeros_like(tilestack.data[0]),
525 positions=ts_aligned.positions,
526 channel=tilestack.channel,
527 czi_fname=tilestack.czi_fname,
528 scene=tilestack.scene,
529 scene_idx=tilestack.scene_idx,
530 )
533def save_tilestack(tilestack: TileStack, output_directory: Path) -> None:
534 """
535 Save each tile in a TileStack as an individual TIFF.
537 Args:
538 tilestack: TileStack containing tile data.
539 output_dir: Base directory for saving.
540 """
541 base = (
542 output_directory
543 / "Tiles"
544 / tilestack.scene
545 / tilestack.czi_fname
546 / tilestack.channel
547 )
548 base.mkdir(parents=True, exist_ok=True)
549 for idx, tile in enumerate(tilestack.data):
550 fname = (
551 f"{tilestack.czi_fname}_{tilestack.scene}_{tilestack.channel}_{idx:04d}.tif"
552 )
553 skimage.io.imsave(base / fname, tile, check_contrast=False)
556def save_stitched_mosaic(mosaic: Mosaic, save_dir: Path, as_gray: bool = False) -> None:
557 """
558 Save a stitched Mosaic to TIFF, optionally as grayscale.
560 Args:
561 mosaic: Mosaic object to save.
562 save_dir: Base output directory.
563 as_gray: If True, convert color mosaics to 8-bit grayscale.
564 """
565 out = save_dir / "Stitched" / mosaic.scene / mosaic.channel
566 out.mkdir(parents=True, exist_ok=True)
567 fname = f"{mosaic.czi_fname}_{mosaic.scene}_{mosaic.channel}_mosaic.tif"
568 data = mosaic.data
569 if as_gray and data.ndim == 3 and data.shape[2] == 3:
570 gray = np.dot(data[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
571 Image.fromarray(gray, mode="L").save(out / fname)
572 else:
573 skimage.io.imsave(out / fname, data, check_contrast=False)
576def normalize_images(img: np.ndarray, flatfield: np.ndarray) -> np.ndarray:
577 """
578 Apply flatfield normalization to an image.
580 Args:
581 img: Input image array.
582 flatfield: Flatfield image array.
584 Returns:
585 Normalized image as uint16.
586 """
587 ff = flatfield.copy().astype(np.float32)
588 ff[ff <= 0] = 1.0
589 norm = img.astype(np.float32) / ff
590 return norm.clip(0, np.iinfo(np.uint16).max).astype(np.uint16)
593def downsample_image(image: np.ndarray, bin_size: int) -> np.ndarray:
594 """
595 Downsample an image by a factor without aliasing.
597 Args:
598 image: Input HxW or HxWxC array.
599 bin_size: Downsampling factor.
601 Returns:
602 Rescaled image.
603 """
604 if image.ndim == 2:
605 new_shape = (image.shape[0] // bin_size, image.shape[1] // bin_size)
606 elif image.ndim == 3:
607 new_shape = (
608 image.shape[0] // bin_size,
609 image.shape[1] // bin_size,
610 image.shape[2],
611 )
612 else:
613 raise ValueError("Unsupported image dimensions.")
614 return skimage.transform.resize(image, new_shape, anti_aliasing=True)
617def prepare_and_save_image(image_array: np.ndarray, file_path: Path) -> None:
618 """
619 Prepare an ndarray and save as TIFF, handling scaling and dtype.
621 Args:
622 image_array: Grayscale or RGB image array.
623 file_path: Destination file path.
624 """
625 arr = image_array
626 if arr.dtype in (np.float32, np.float64):
627 arr = np.clip(arr, 0.0, 1.0)
628 arr = (arr * 255).astype(np.uint8)
629 elif arr.dtype != np.uint8:
630 arr = arr.astype(np.uint8)
631 if arr.ndim == 2:
632 img = Image.fromarray(arr, mode="L")
633 elif arr.ndim == 3 and arr.shape[2] == 3:
634 img = Image.fromarray(arr, mode="RGB")
635 else:
636 raise ValueError("Unsupported image format.")
637 img.save(file_path, format="TIFF")
640def ezshow(mosaics: List[Mosaic]) -> None:
641 """
642 Display a list of mosaics using matplotlib.
644 Args:
645 mosaics: List of Mosaic objects.
646 """
647 for m in mosaics:
648 img = (m.data / np.max(m.data) * 255).astype(np.uint8)
649 fig, ax = plt.subplots()
650 ax.imshow(img, cmap="gray" if img.ndim == 2 else None)
651 ax.set_title(f"{m.czi_fname} | Scene: {m.scene} | Channel: {m.channel}")
652 plt.show()