Coverage for src/precon3d/stitcher.py: 20%
158 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"""
4stitcher.py
6A unified CLI for stitching ZEISS CZI mosaics and tile-region arrays.
8Features:
9 * Save “simple” downscaled mosaics (color or grayscale).
10 * Stitch per-channel mosaics in parallel using a reference channel.
11 * Optional flatfield normalization for mosaics.
12 * Handle non-mosaic CZI files as tile-region arrays.
13 * Commands for checking tile arrays and processing directories.
14"""
16import time
17from pathlib import Path
18from typing import List
20import typer
21import numpy as np
22import skimage.io
23import imageio
24from PIL import Image
25from rich.progress import Progress, SpinnerColumn, TextColumn
26from joblib import Parallel, delayed
27from aicspylibczi import CziFile
29# Local imports
30import precon3d.mosaic_utils as mosaic_utils
31import precon3d.utility as ut
32import precon3d.factory as factory
33from precon3d._my_typer_cli import CustomCLIGroup, CustomCLICommand
34from precon3d.custom_types import (
35 StitchingParams,
36 GeneralAttrs,
37 StitchingAttrs,
38 FijiAttrs,
39 NormalizationAttrs,
40 TileStack,
41 Mosaic,
42 TilePosition,
43)
45app = typer.Typer(
46 cls=CustomCLIGroup,
47 no_args_is_help=True,
48 help="Stitch tiles and mosaics from ZEISS CZI files.",
49)
52def save_simple_color_mosaic(
53 czi_file: Path,
54 general_attrs: GeneralAttrs,
55 stitching_attrs: StitchingAttrs,
56 fiji_attrs: FijiAttrs,
57) -> None:
58 """
59 Save downscaled, per-channel color mosaics for each scene.
61 Args:
62 czi_file (Path): Path to the CZI file.
63 general_attrs (GeneralAttrs): General pipeline settings.
64 stitching_attrs (StitchingAttrs): Stitching parameters.
65 fiji_attrs (FijiAttrs): Fiji application settings.
66 """
67 tilestack, fname = mosaic_utils.ai_tilestack(czi_file)
68 bin_size = stitching_attrs.simple_mosaic_downscale_bin_size
70 for scene_idx in tilestack.scenes:
71 stacks = mosaic_utils.simple_tilestack(tilestack, scene_idx, fname)
72 for ts in stacks:
73 # stitch the tiles using Fiji (per-channel)
74 stitched = mosaic_utils.stitch_tilestack(
75 ts, fiji_attrs.fiji_app, cleanup=True, tileconfig=ts.positions
76 )
77 # prepare output directory and filename
78 out_dir = (
79 general_attrs.output_directory
80 / "Simple_Mosaic"
81 / stitched.scene
82 / stitched.channel
83 )
84 out_dir.mkdir(parents=True, exist_ok=True)
85 fname_out = (
86 f"{stitched.czi_fname}_{stitched.scene}_"
87 f"{stitched.channel}_binsz{bin_size:02d}_mosaic.tif"
88 )
89 # downscale
90 img = stitched.data.squeeze()
91 new_shape = (
92 img.shape[0] // bin_size,
93 img.shape[1] // bin_size,
94 img.shape[2],
95 )
96 img_resized = skimage.transform.resize(
97 img, new_shape, anti_aliasing=True
98 )
99 # save 8-bit TIFF
100 skimage.io.imsave(
101 out_dir / fname_out,
102 ut.convert_gray_to_8bit(img_resized.squeeze()),
103 check_contrast=False,
104 )
105 typer.echo(f'✓ Simple color mosaics from "{fname}" saved.')
108def save_simple_mosaic(
109 czi_file: Path,
110 general_attrs: GeneralAttrs,
111 stitching_attrs: StitchingAttrs,
112 fiji_attrs: FijiAttrs,
113) -> None:
114 """
115 Save downscaled mosaics (color or grayscale) generated
116 directly from acquisition stage positions.
118 Args:
119 czi_file (Path): Path to the CZI file.
120 general_attrs (GeneralAttrs): General pipeline settings.
121 stitching_attrs (StitchingAttrs): Stitching parameters.
122 fiji_attrs (FijiAttrs): Fiji application settings.
123 """
124 ai_obj, _ = mosaic_utils.ai_tilestack(czi_file)
125 dims = dict(ai_obj.dims)
126 if dims.get("S", False):
127 save_simple_color_mosaic(
128 czi_file, general_attrs, stitching_attrs, fiji_attrs
129 )
130 return
132 # grayscale path
133 czi_mos, fname = mosaic_utils.ai_mosaic(czi_file)
134 bin_size = stitching_attrs.simple_mosaic_downscale_bin_size
136 for scene_idx in czi_mos.scenes:
137 czi_mos.set_scene(scene_idx)
138 channels = dims.get("C", 1)
139 channel_mosaics = []
140 for ch in range(channels):
141 arr = czi_mos.get_image_data("MYXS", C=ch).squeeze()
142 positions = czi_mos.reader.get_mosaic_tile_positions()
143 tiles = [TilePosition(y=y, x=x) for y, x in positions]
144 channel_mosaics.append(
145 TileStack(
146 data=arr,
147 positions=tiles,
148 czi_fname=fname,
149 channel=czi_mos.channel_names[ch],
150 scene=czi_mos.current_scene,
151 scene_idx=czi_mos.current_scene_index,
152 )
153 )
154 # downscale & save each channel
155 for ts in channel_mosaics:
156 out_dir = (
157 general_attrs.output_directory
158 / "Simple_Mosaic"
159 / ts.scene
160 / ts.channel
161 )
162 out_dir.mkdir(parents=True, exist_ok=True)
163 fname_out = (
164 f"{ts.czi_fname}_{ts.scene}_{ts.channel}"
165 f"_binsz{bin_size:02d}_mosaic.tif"
166 )
167 img = ts.data.squeeze()
168 pil = Image.fromarray(img)
169 pil = pil.resize(
170 (img.shape[1] // bin_size, img.shape[0] // bin_size),
171 Image.Resampling.LANCZOS,
172 )
173 skimage.io.imsave(
174 out_dir / fname_out,
175 ut.convert_gray_to_8bit(np.array(pil).squeeze()),
176 check_contrast=False,
177 )
178 typer.echo(f'✓ Simple grayscale mosaics from "{fname}" saved.')
181def normalize_simple_tilestacks(
182 stacks: List[TileStack],
183 norm_attrs: NormalizationAttrs,
184) -> List[TileStack]:
185 """
186 Normalize each TileStack by its channel-specific flatfield.
188 Args:
189 stacks (List[TileStack]): Tile stacks to normalize.
190 norm_attrs (NormalizationAttrs): Flatfield settings.
192 Returns:
193 List[TileStack]: New TileStacks with normalized data.
195 Raises:
196 ValueError: If a required flatfield file is missing.
197 """
198 parent = Path(norm_attrs.channel_flatfields_parent_directory)
199 normalized = []
201 for ts in stacks:
202 ff_name = norm_attrs.channel_flatfields_filename.get(ts.channel)
203 if ff_name is None:
204 raise ValueError(f"No flatfield entry for channel '{ts.channel}'")
205 ff_path = parent / ff_name
206 if not ff_path.exists():
207 raise ValueError(f"Flatfield file '{ff_path}' not found")
208 ff_img = imageio.v3.imread(ff_path).squeeze()
209 if ff_img.ndim == 2:
210 ff_img = ff_img[..., None]
212 data = ts.data
213 norm_tiles = np.zeros_like(data, dtype=np.uint16)
214 for i, tile in enumerate(data):
215 norm_tiles[i] = mosaic_utils.normalize_images(tile, ff_img)
216 normalized.append(
217 TileStack(
218 data=norm_tiles,
219 positions=ts.positions,
220 channel=ts.channel,
221 czi_fname=ts.czi_fname,
222 scene=ts.scene,
223 scene_idx=ts.scene_idx,
224 )
225 )
226 return normalized
229def stitch_simple_tilestacks_with_reference_channel(
230 stacks: List[TileStack],
231 fiji_attrs: FijiAttrs,
232 ref_channel: str,
233) -> List[Mosaic]:
234 """
235 Stitch a list of TileStacks, using one channel as reference.
237 The reference channel is stitched first, its tile positions
238 are then applied to other channels to guarantee alignment.
240 Args:
241 stacks (List[TileStack]): Per-channel TileStacks.
242 fiji_attrs (FijiAttrs): Fiji application settings.
243 ref_channel (str): Channel name to use as reference.
245 Returns:
246 List[Mosaic]: Stitched mosaics for all channels.
247 """
248 ref_ts = next((ts for ts in stacks if ts.channel == ref_channel), None)
249 if ref_ts is None:
250 raise ValueError(f"Reference channel '{ref_channel}' not found")
252 ref_mos = mosaic_utils.stitch_tilestack(
253 ref_ts, fiji_attrs.fiji_app, cleanup=True
254 )
255 out = [ref_mos]
256 others = [ts for ts in stacks if ts.channel != ref_channel]
258 for ts in others:
259 m = mosaic_utils.stitch_tilestack(
260 ts,
261 fiji_attrs.fiji_app,
262 cleanup=True,
263 tileconfig=ref_mos.positions,
264 )
265 out.append(m)
266 return out
269def stitch_and_save_mosaic(czi_file: Path, params: StitchingParams) -> None:
270 """
271 Full pipeline: per-scene tilestack → normalize → stitch → save.
273 Args:
274 czi_file (Path): Path to a mosaic CZI file.
275 params (StitchingParams): All pipeline parameters.
276 """
277 tilestack, fname = mosaic_utils.ai_tilestack(czi_file)
279 for scene_idx in tilestack.scenes:
280 stacks = mosaic_utils.simple_tilestack(tilestack, scene_idx, fname)
282 if params.normalization_attrs.use_flatfield:
283 stacks = normalize_simple_tilestacks(
284 stacks, params.normalization_attrs
285 )
287 mosaics = stitch_simple_tilestacks_with_reference_channel(
288 stacks,
289 params.fiji_attrs,
290 params.stitching_attrs.stitch_reference_channel,
291 )
292 # save
293 for m in mosaics:
294 # color check
295 out_dir = (
296 params.general_attrs.output_directory
297 / "Stitched"
298 / m.scene
299 / m.channel
300 )
301 out_dir.mkdir(parents=True, exist_ok=True)
302 fname_out = f"{m.czi_fname}_{m.scene}_{m.channel}_mosaic.tif"
303 img = m.data
304 if img.shape[-1] == 3:
305 # RGB: convert BGR→RGB and gamma‐correct
306 rgb = img[:, :, ::-1]
307 rgb = skimage.exposure.adjust_gamma(rgb, 0.5)
308 skimage.io.imsave(
309 out_dir / fname_out, rgb, check_contrast=False
310 )
311 else:
312 skimage.io.imsave(
313 out_dir / fname_out,
314 ut.convert_gray_to_8bit(img.squeeze()),
315 check_contrast=False,
316 )
319def stitch_save_tilearray_by_idx(
320 czi_file: Path,
321 params: StitchingParams,
322 grouping_idx: int,
323) -> None:
324 """
325 Stitch and save a tile‐region array CZI by grouping index.
327 Args:
328 czi_file (Path): Path to a tile‐region CZI file.
329 params (StitchingParams): Pipeline parameters.
330 grouping_idx (int): Tile‐array group index.
331 """
332 stacks = mosaic_utils.retrieve_tile_region_array_data(czi_file)
333 grouped = mosaic_utils.group_tile_arrays(stacks, idx=grouping_idx)
334 for ts in grouped:
335 mos = mosaic_utils.stitch_tiles_from_positions(ts)
336 mosaic_utils.save_stitched_mosaic(
337 mos, params.general_attrs.output_directory
338 )
341def start_stitcher_parallel(
342 params: StitchingParams,
343 ncores: int = 4,
344 grouping_idx: int = 0,
345) -> None:
346 """
347 Process all CZI files in input directory in parallel.
349 Depending on `stitch_tiles` and whether the file is a mosaic,
350 either runs full‐mosaic stitching or tile‐region array stitching.
352 Args:
353 params (StitchingParams): Pipeline parameters.
354 ncores (int): Number of parallel jobs.
355 grouping_idx (int): Group index for tile‐region arrays.
356 """
357 ext = params.general_attrs.file_extension
358 if ext != ".czi":
359 raise ValueError("Only .czi is supported")
361 inp = params.general_attrs.input_directory
362 all_files = ut.sorted_files(inp, ext)
363 out_stitched = params.general_attrs.output_directory / "Stitched"
365 if out_stitched.exists():
366 done = ut.sorted_files(params.general_attrs.output_directory, ".tif")
367 todo = ut.remove_matching_filenames(all_files, done)
368 else:
369 todo = all_files
371 if params.stitching_attrs.stitch_tiles:
372 first = todo[0]
373 if CziFile(first).is_mosaic():
374 # full mosaic path
375 Parallel(n_jobs=ncores)(
376 delayed(stitch_and_save_mosaic)(f, params) for f in todo
377 )
378 else:
379 # tile‐region arrays
380 Parallel(n_jobs=ncores)(
381 delayed(stitch_save_tilearray_by_idx)(f, params, grouping_idx)
382 for f in todo
383 )
386@app.command(
387 cls=CustomCLICommand,
388 name="process_images",
389 help="Stitch all CZI in parallel per config.",
390)
391def process_images(
392 configfile: Path = typer.Argument(..., exists=True, help="YAML config."),
393 ncores: int = typer.Option(4, help="Number of parallel jobs."),
394 grouping_idx: int = typer.Option(
395 0, help="Tile-array group index for non-mosaics."
396 ),
397) -> None:
398 """
399 Read a YAML config and process all CZI files in the input directory.
401 Example:
402 python stitcher.py process_images config.yml --ncores 8
403 """
404 config = ut.read_config(configfile)
405 params = factory.create_stitching_parameters(config)
406 ut.current_date_and_time()
407 typer.secho(f"*** Starting processing: {configfile.name} ***", fg="green")
408 t0 = time.perf_counter()
409 start_stitcher_parallel(params, ncores=ncores, grouping_idx=grouping_idx)
410 dt = time.perf_counter() - t0
411 typer.secho(f"Total time: {dt:.2f}s ({dt/60:.2f}min)", fg="green")
414@app.command(
415 cls=CustomCLICommand,
416 name="stitch_czi",
417 help="Stitch a single CZI with progress spinner.",
418)
419def stitch_czi(
420 czifile: Path = typer.Argument(..., exists=True, help="CZI to stitch."),
421 configfile: Path = typer.Argument(..., exists=True, help="YAML config."),
422) -> None:
423 """
424 Stitch one CZI (must be a mosaic) using the provided config.
426 Example:
427 python stitcher.py stitch_czi sample.czi config.yml
428 """
429 czi = czifile
430 cfg = ut.read_config(configfile)
431 params = factory.create_stitching_parameters(cfg)
432 ut.current_date_and_time()
433 typer.secho(f"*** Stitching {czi.name} ***", fg="cyan")
434 t0 = time.perf_counter()
435 with Progress(
436 SpinnerColumn(),
437 TextColumn("[progress.description]{task.description}"),
438 transient=True,
439 ) as prog:
440 task = prog.add_task("Stitching...", total=None)
441 stitch_and_save_mosaic(czi, params)
442 prog.update(task, completed=True)
443 dt = time.perf_counter() - t0
444 typer.secho(f"Done in {dt:.2f}s ({dt/60:.2f}min)", fg="cyan")
447@app.command(
448 cls=CustomCLICommand,
449 name="check_tilearray",
450 help="Visualize tile-region arrays in a CZI.",
451)
452def check_tilearray(
453 czi_path: Path = typer.Argument(..., exists=True, help="CZI to inspect."),
454) -> None:
455 """
456 Load tile-region array data and display a quick overview.
458 Example:
459 python stitcher.py check_tilearray sample.czi
460 """
461 stacks = mosaic_utils.tile_region_array_data(czi_path)
462 flat = mosaic_utils.flatten_list(stacks)
463 mosaic_utils.visualize_tilestacks(flat)
466if __name__ == "__main__":
467 app()