Coverage for src/precon3d/shading_correction.py: 16%
271 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 03:50 +0000
1# shading correction
2import os
3from sys import platform
4import shutil
5import skimage
6import skimage.io as skio
7import imageio
8from pathlib import Path
9import numpy as np
10import time
11import random
12from dataclasses import dataclass, field
13from joblib import Parallel, delayed
16from numpy.typing import NDArray
17from typing import List, Dict, Any, NamedTuple
18from scipy.stats import wasserstein_distance, chisquare
19from sklearn.decomposition import NMF
20from subprocess import Popen, DEVNULL, PIPE
21import subprocess
22import typer
24# import dask.array as da
25# import dask_image.imread
27# from PIL import Image
28from tqdm import tqdm
30# local libraries
31import precon3d.utility as ut
32import precon3d.mosaic_utils
33import precon3d.factory as gen_types
35# pylint: disable=wildcard-import
36from precon3d.custom_types import *
37from precon3d._my_typer_cli import CustomCLIGroup, CustomCLICommand
39# CLI
40app = typer.Typer(
41 cls=CustomCLIGroup,
42 no_args_is_help=True,
43 short_help="Estimate shading correction from tiles",
44)
47def subimages_similar_by_ssim(
48 ref_image, found_image, n_subimages, ssim_thresh
49):
50 """Compare subimages are also similar
52 divide into equal sized sub images and compare they are all similar
53 """
54 height, width = ref_image.shape[:2]
55 subimg_height = height // n_subimages
56 subimg_width = width // n_subimages
58 # Pre-calculate the slicing indices to avoid doing it in the loop
59 slice_indices = [
60 (y, x)
61 for y in range(0, height, subimg_height)
62 for x in range(0, width, subimg_width)
63 if (y + subimg_height <= height and x + subimg_width <= width)
64 ]
66 ssim_scores = []
67 for y, x in slice_indices:
68 if found_image.shape[-1] == 3: # COLOR
69 ref_subimg = ref_image[
70 y : y + subimg_height, x : x + subimg_width, :
71 ].squeeze()
72 found_subimg = found_image[
73 y : y + subimg_height, x : x + subimg_width, :
74 ].squeeze()
76 # Calculate SSIM for the current pair of subimages
77 score, _ = skimage.metrics.structural_similarity(
78 ref_subimg,
79 found_subimg,
80 data_range=65535,
81 full=True,
82 channel_axis=-1,
83 )
84 else:
85 ref_subimg = ref_image[
86 y : y + subimg_height, x : x + subimg_width
87 ].squeeze()
88 found_subimg = found_image[
89 y : y + subimg_height, x : x + subimg_width
90 ].squeeze()
92 # Calculate SSIM for the current pair of subimages
93 score, _ = skimage.metrics.structural_similarity(
94 ref_subimg,
95 found_subimg,
96 data_range=65535,
97 full=True,
98 )
99 ssim_scores.append(score)
101 # Check if all SSIM scores are above the threshold
102 return all(score > ssim_thresh for score in ssim_scores)
105def downselect_tiles(
106 reorganized_tile_dir: Path,
107 ref_ch: str,
108 ref_img_idx: int = 0,
109 ssim_thresh: float = 0.9,
110 n_subimages: int = 10,
111) -> bool:
112 """Downselect tiles organized already by channels"""
114 # get channel subdirectories
115 channel_subdirs = [
116 ch for ch in reorganized_tile_dir.iterdir() if ch.is_dir()
117 ]
119 # get the file paths
120 all_file_paths = {
121 subdir.name: ut.sorted_files(subdir, ".tif")
122 for subdir in channel_subdirs
123 }
125 print(f"Found {len(channel_subdirs)} channels")
126 # Dictionary to hold the count of file paths for each subdir
127 file_path_counts = {
128 subdir: len(paths) for subdir, paths in all_file_paths.items()
129 }
131 # Printing the number of file paths for each subdir
132 for subdir, count in file_path_counts.items():
133 print(f"\t{subdir}: {count} file(s)")
135 # Get an iterator of all list lengths
136 lengths = (len(paths) for paths in all_file_paths.values())
138 # Get the length of the first list to compare against others
139 first_length = next(
140 lengths, None
141 ) # Returns None if no elements are present
143 # Check if all lengths are the same as the first length
144 all_same_length = all(length == first_length for length in lengths)
145 if all_same_length:
146 print("all folders have the same number of files")
147 else:
148 # Raise an exception with a message indicating the problem
149 # raise ValueError('Folders have different numbers of files. Each folder must contain the same number of files.')
151 file_counts = {
152 subdir: len(paths) for subdir, paths in all_file_paths.items()
153 }
154 error_message = (
155 "Folders have different numbers of files: "
156 + ", ".join(f"{k}: {v} files" for k, v in file_counts.items())
157 )
158 raise ValueError(error_message)
160 ref_filepath = all_file_paths.get(ref_ch)[ref_img_idx]
161 ref_image = skio.imread(ref_filepath)
163 print(
164 f"Downselecting tiles in {reorganized_tile_dir} by {ref_ch} channel\nUsing {ref_filepath} as reference"
165 )
167 for each_ch_subdir in channel_subdirs:
168 all_file_paths[each_ch_subdir.name] = ut.sorted_files(
169 each_ch_subdir, "tif"
170 ) # each_ch_subdir.iterdir()# sort naturally for the files in the directories
172 # downselect based on the reference channel
173 for idx, each_image in enumerate(all_file_paths.get(ref_ch)):
174 found_image = skio.imread(each_image)
176 if found_image.shape[-1] == 3: # COLOR
177 ssim_score = skimage.metrics.structural_similarity(
178 ref_image.squeeze(),
179 found_image.squeeze(),
180 data_range=65535,
181 full=True,
182 channel_axis=-1,
183 )
184 else:
185 ssim_score = skimage.metrics.structural_similarity(
186 ref_image.squeeze(),
187 found_image.squeeze(),
188 data_range=65535,
189 full=True,
190 )
192 # remove if not similar
193 if (
194 ssim_score[0] <= ssim_thresh
195 ): # not similar based on the ssim for subimages
196 print(
197 f"DIFFERENT: SSIM score = {ssim_score[0]} ({each_image.stem})"
198 )
199 # delete specific file from all dirs
200 for subdir_name, file_path in all_file_paths.items():
201 file_to_delete = file_path[idx] # use index to specify
203 print(f"\t{file_to_delete.stem} deleted from {subdir_name}")
205 file_to_delete.unlink()
207 else: # overall, images look similar, then divide into subimages and compare subimages
208 print(
209 f"SIMILAR, check subimages: SSIM score = {ssim_score[0]} ({each_image.stem})"
210 )
212 subimages_similar = subimages_similar_by_ssim(
213 ref_image, found_image, n_subimages, ssim_thresh
214 )
216 if subimages_similar is False:
217 # not similar based on the ssim for subtiles
219 print(
220 f"\tFound differences when tiles divided into {n_subimages*n_subimages} smaller pieces"
221 )
223 # breakpoint()
225 # delete specific file from all dirs
226 for subdir_name, file_path in all_file_paths.items():
227 file_to_delete = file_path[idx] # use index to specify
229 print(
230 f"\t{file_to_delete.stem} deleted from {subdir_name}"
231 )
232 # print(f"\t\t{ssim_list}")
234 file_to_delete.unlink()
237@app.command(
238 cls=CustomCLICommand,
239 name="downselect_tiles_in_directory",
240 short_help="Use a SSIM threshold to downselect tiles",
241)
242def downselect_tiles_in_directory(
243 reorganized_tile_dir: Path,
244 ref_img_idx: int = 0,
245 ssim_thresh: float = 0.9,
246 n_subimages: int = 10,
247) -> bool:
248 """Downselect tiles organized already by channels"""
250 # get the file paths
252 all_file_paths = ut.sorted_files(
253 reorganized_tile_dir, ".tif"
254 ) # each_ch_subdir.iterdir()# sort naturally for the files in the directories
255 ref_image = skio.imread(all_file_paths[ref_img_idx])
257 # downselect based on the reference channel
258 for each_image in all_file_paths:
259 found_image = skio.imread(each_image)
261 if found_image.shape[-1] == 3: # COLOR
262 ssim_score = skimage.metrics.structural_similarity(
263 ref_image.squeeze(),
264 found_image.squeeze(),
265 data_range=65535,
266 full=True,
267 channel_axis=-1,
268 )
269 else:
270 ssim_score = skimage.metrics.structural_similarity(
271 ref_image.squeeze(),
272 found_image.squeeze(),
273 data_range=65535,
274 full=True,
275 )
277 # remove if not similar
278 if (
279 ssim_score[0] <= ssim_thresh
280 ): # not similar based on the ssim for subimages
281 print(
282 f"DIFFERENT: SSIM score = {ssim_score[0]} ({each_image.stem})"
283 )
284 # delete specific file
285 each_image.unlink()
287 else: # overall, images look similar, then divide into subimages and compare subimages
288 print(
289 f"SIMILAR, check subimages: SSIM score = {ssim_score[0]} ({each_image.stem})"
290 )
292 subimages_similar = subimages_similar_by_ssim(
293 ref_image, found_image, n_subimages, ssim_thresh
294 )
296 if subimages_similar is False:
297 # not similar based on the ssim for subtiles
299 print(
300 f"\tFound differences when tiles divided into {n_subimages*n_subimages} smaller pieces"
301 )
303 # delete specific file
304 each_image.unlink()
307def downselect_tiles_in_directory_parallel(
308 reorganized_tile_dir: Path,
309 ref_img_idx: int = 0,
310 ssim_thresh: float = 0.9,
311 n_subimages: int = 10,
312 n_cores: int = 8,
313) -> bool:
314 """Downselect tiles organized already by channels"""
316 # get the file paths
318 all_file_paths = ut.sorted_files(
319 reorganized_tile_dir, ".tif"
320 ) # each_ch_subdir.iterdir()# sort naturally for the files in the directories
321 ref_image = skio.imread(all_file_paths[ref_img_idx])
323 # downselect based on the reference channel, parallel
324 def _compare_image_with_ssim(comparison_image: Path):
325 """tocome"""
327 found_image = skio.imread(comparison_image)
329 if found_image.shape[-1] == 3: # COLOR
330 ssim_score = skimage.metrics.structural_similarity(
331 ref_image.squeeze(),
332 found_image.squeeze(),
333 data_range=65535,
334 full=True,
335 channel_axis=-1,
336 )
337 else:
338 ssim_score = skimage.metrics.structural_similarity(
339 ref_image.squeeze(),
340 found_image.squeeze(),
341 data_range=65535,
342 full=True,
343 )
345 # remove if not similar
346 if (
347 ssim_score[0] <= ssim_thresh
348 ): # not similar based on the ssim for subimages
349 print(
350 f"DIFFERENT: SSIM score = {ssim_score[0]} ({comparison_image.stem})"
351 )
352 # delete specific file
353 comparison_image.unlink()
355 else: # overall, images look similar, then divide into subimages and compare subimages
356 print(
357 f"SIMILAR, check subimages: SSIM score = {ssim_score[0]} ({each_image.stem})"
358 )
360 subimages_similar = subimages_similar_by_ssim(
361 ref_image, found_image, n_subimages, ssim_thresh
362 )
364 if subimages_similar is False:
365 # not similar based on the ssim for subtiles
367 print(
368 f"\tFound differences when tiles divided into {n_subimages*n_subimages} smaller pieces"
369 )
371 # delete specific file
372 comparison_image.unlink()
374 Parallel(n_jobs=n_cores)(
375 delayed(_compare_image_with_ssim)(each_image)
376 for each_image in all_file_paths
377 )
380def downselect_tiles_using_reference_channel(
381 reorganized_tile_dir: Path,
382 ref_ch: str,
383 ref_image: Path,
384 ssim_thresh: float = 0.9,
385 n_subimages: int = 10,
386) -> bool:
387 """Downselect tiles organized already by channels"""
389 ref_image = skio.imread(ref_image)
391 # get channel subdirectories
392 channel_subdirs = [
393 ch for ch in reorganized_tile_dir.iterdir() if ch.is_dir()
394 ]
396 # get the file paths
397 all_file_paths = {subdir.name: [] for subdir in channel_subdirs}
398 for each_ch_subdir in channel_subdirs:
399 all_file_paths[each_ch_subdir.name] = ut.sorted_files(
400 each_ch_subdir, "tif"
401 ) # each_ch_subdir.iterdir()# sort naturally for the files in the directories
403 # check number of files are the same
404 n_files = [len(file_paths) for file_paths in all_file_paths.values()]
405 if len(set(n_files)) == 1:
406 print(
407 f"The channel dirs: {list(all_file_paths.keys())}, all contain {list(set(n_files))} files\n"
408 )
409 else:
410 print(
411 f"The channel dirs: {list(all_file_paths.keys())}, contain {n_files} files\n"
412 )
413 return
415 # downselect based on the reference channel
416 for idx, each_image in enumerate(all_file_paths.get(ref_ch)):
417 found_image = skio.imread(each_image)
419 if found_image.shape[-1] == 3: # COLOR
420 ssim_score = skimage.metrics.structural_similarity(
421 ref_image.squeeze(),
422 found_image.squeeze(),
423 data_range=65535,
424 full=True,
425 channel_axis=-1,
426 )
427 else:
428 ssim_score = skimage.metrics.structural_similarity(
429 ref_image.squeeze(),
430 found_image.squeeze(),
431 data_range=65535,
432 full=True,
433 )
435 # remove if not similar
436 if (
437 ssim_score[0] <= ssim_thresh
438 ): # not similar based on the ssim for subimages
439 print(
440 f"DIFFERENT: SSIM score = {ssim_score[0]} ({each_image.stem})"
441 )
442 # delete specific file from all dirs
443 for subdir_name, file_path in all_file_paths.items():
444 file_to_delete = file_path[idx] # use index to specify
446 print(f"\t{file_to_delete.stem} deleted from {subdir_name}")
448 file_to_delete.unlink()
450 else: # overall, images look similar, then divide into subimages and compare subimages
451 print(
452 f"SIMILAR, check subimages: SSIM score = {ssim_score[0]} ({each_image.stem})"
453 )
455 if n_subimages == 1:
456 continue
458 subimages_similar = subimages_similar_by_ssim(
459 ref_image, found_image, n_subimages, ssim_thresh
460 )
462 if subimages_similar is False:
463 # not similar based on the ssim for subtiles
465 print(
466 f"\tFound differences when tiles divided into {n_subimages*n_subimages} smaller pieces"
467 )
469 # breakpoint()
471 # delete specific file from all dirs
472 for subdir_name, file_path in all_file_paths.items():
473 file_to_delete = file_path[idx] # use index to specify
475 print(
476 f"\t{file_to_delete.stem} deleted from {subdir_name}"
477 )
478 # print(f"\t\t{ssim_list}")
480 file_to_delete.unlink()
483# def nmf_on_images(images: da.Array) -> np.ndarray:
484# """
485# Apply NMF to a set of images and return the flatfield correction.
487# Parameters:
488# - images: A Dask array of images.
490# Returns:
491# - A NumPy array representing the flatfield correction.
492# """
493# # Reshape images for NMF: (n_samples, n_features)
494# reshaped_images = images.reshape((images.shape[0], -1))
496# # Convert to NumPy array (compute the Dask array)
497# reshaped_images_np = reshaped_images.compute()
499# # Apply NMF
500# model = NMF(n_components=1, init="random", random_state=0)
501# _ = model.fit_transform(reshaped_images_np)
502# h = model.components_
504# # Assume the first image's shape is representative
505# flatfield = h.reshape(images.shape[1:])
507# # Normalize the flatfield
508# flatfield_mean = np.mean(flatfield)
509# if flatfield_mean != 0:
510# flatfield /= flatfield_mean
512# return flatfield.astype(np.float32)
515# def process_in_parts(shading_tiles_dir: Path, n_parts: int) -> np.ndarray:
516# """
517# Split the reading and processing of images into n_parts, apply NMF to each part,
518# and average the results to get the final flatfield correction.
520# Parameters:
521# - shading_tiles_dir: Path to the directory containing the shading images.
522# - n_parts: Number of parts to split the dataset into.
524# Returns:
525# - A NumPy array representing the averaged flatfield correction.
526# """
527# # Load all image paths
528# # image_paths = list(shading_tiles_dir.glob('*.tif'))
529# image_paths = ut.sorted_files(shading_tiles_dir, ".tif")
530# total_images = len(image_paths)
531# images_per_part = total_images // n_parts
533# flatfields = []
535# for part in range(n_parts):
536# print(f"Part {part+1} of {n_parts}")
538# start_idx = part * images_per_part
539# end_idx = (
540# (part + 1) * images_per_part
541# if part < n_parts - 1
542# else total_images
543# )
545# # Use Dask to lazily load a subset of images
546# # images = dask_image.imread.imread(image_paths[start_idx:end_idx])
547# images = [
548# dask_image.imread.imread(each_path)[None, ...]
549# for each_path in image_paths[start_idx:end_idx]
550# ]
551# image_stack = da.concatenate(images, axis=0)
552# # Apply NMF and get the flatfield correction for this part
553# flatfield = nmf_on_images(image_stack)
554# flatfields.append(flatfield)
556# return flatfields
557# # # Average the flatfield corrections from all parts
558# # averaged_flatfield = np.mean(flatfields, axis=0)
560# # return averaged_flatfield
563# def estimate_shading_correction_with_dask(shading_tiles_dir: Path) -> NDArray:
564# """
565# Estimate the shading correction for a set of large images using Non-negative Matrix Factorization (NMF),
566# with optimizations for memory constraints by using Dask arrays for lazy loading and processing.
568# Parameters
569# ----------
570# shading_tiles_dir : Path
571# The directory path where the large shading images are stored. These images are used to compute the shading correction.
573# Returns
574# -------
575# NDArray
576# A 2D numpy array representing the estimated flatfield correction image. This image is used to correct
577# other images for shading effects, normalized so its mean is centered around 1.
579# Notes
580# -----
581# - This function is optimized for handling large images that do not fit into memory by using Dask for lazy evaluation.
582# - The function assumes that all shading images are of the same dimensions and are stored in '.tif' format.
583# - The flatfield correction is derived from the NMF decomposition of the shading images, specifically from the H matrix,
584# and is normalized to have a mean of 1 for consistent lighting correction.
585# """
587# # Use dask_image to lazily load the images as a Dask array
588# images = dask_image.imread.imread(str(shading_tiles_dir / '*.tif'))
590# # Since Dask operates lazily, computations are only triggered upon calling compute().
591# # Here, we reshape the images into a 2D array where each row is a flattened image.
592# # Note: This step might need adjustments based on the actual memory constraints and image sizes.
593# reshaped_images = images.reshape((images.shape[0], -1)).compute()
595# # Apply NMF on the reshaped images to find the flatfield correction
596# # Note: NMF is performed on the in-memory NumPy array since sklearn's NMF does not directly support Dask arrays.
597# model = NMF(n_components=1, init='random', random_state=0)
598# W = model.fit_transform(reshaped_images)
599# H = model.components_
601# # Reshape the H matrix to get the flatfield correction image
602# # The first image's shape is used as a reference for reshaping
603# flatfield_shape = images.shape[1:] # Ignoring the number of images, just getting one image's shape
604# flatfield = H.reshape(flatfield_shape)
606# # Normalize the flatfield so its mean is centered around 1
607# flatfield_mean = np.mean(flatfield)
608# if flatfield_mean != 0:
609# flatfield /= flatfield_mean
611# print(f'Estimated the flatfield ({flatfield.shape[0]}x{flatfield.shape[1]} pixels) for normalization with non-negative matrix factorization.')
613# return flatfield
616@app.command(
617 cls=CustomCLICommand,
618 name="estimate_shading_correction",
619 short_help="Estimate flatfield with NNMF",
620)
621def estimate_shading_correction(
622 shading_tiles_dir: Path, n_images_limit: int = 800
623):
624 """
625 Estimate the shading correction for a set of images using Non-negative Matrix Factorization (NMF).
627 This function takes a directory containing shading images, loads them, and applies NMF to estimate
628 a flatfield correction image. The flatfield image is normalized so that its mean is centered around 1.
629 This normalization is crucial for subsequent image processing steps, ensuring that the corrected images
630 have uniform lighting conditions.
632 Parameters
633 ----------
634 shading_tiles_dir : Path
635 The directory path where the shading images are stored. These images are used to compute the shading correction.
637 Returns
638 -------
639 NDArray
640 A 2D numpy array representing the estimated flatfield correction image. This image is used to correct
641 other images for shading effects.
643 Notes
644 -----
645 - The function assumes that all shading images are of the same dimensions.
646 - The shading images should be in '.tif' format.
647 - The function uses Non-negative Matrix Factorization (NMF) to decompose the shading images into a basis
648 component (W) and coefficient matrix (H). The flatfield correction is derived from these components.
650 Examples
651 --------
652 >>> shading_correction = estimate_shading_correction(Path('/path/to/shading/images'))
653 >>> print(shading_correction.shape)
654 (1024, 1024) # Example output, the actual size depends on the input images.
655 """
657 # Retrieve a sorted list of shading image files from the specified directory
658 shading_images = ut.sorted_files(shading_tiles_dir, ".tif")
660 if len(shading_images) > n_images_limit:
661 print(
662 f"Woah, found {len(shading_images)} images in {shading_tiles_dir}, selecting random subset of {n_images_limit} images\n"
663 )
664 shading_images = random.sample(shading_images, n_images_limit)
665 else:
666 print(f"Found {len(shading_images)} images in {shading_tiles_dir}\n")
668 # Load the shading images into a list, using imageio for image reading
669 images = [
670 imageio.v3.imread(each_shading_image)
671 for each_shading_image in tqdm(shading_images)
672 ]
674 ff_shape = images[0].shape
675 # Reshape the loaded images into a 2D array where each row represents a flattened image
676 reshaped_images = np.array([img.flatten() for img in images])
678 del images
680 # Initialize and fit the NMF model to the reshaped images
681 # n_components=1 indicates we are reducing the image set to 1 component
682 model = NMF(n_components=1, init="random", random_state=0)
683 W = model.fit_transform(reshaped_images)
685 del reshaped_images
687 # Basis matrix
688 H = model.components_ # Coefficient matrix
690 # Reshape the H matrix to the original image shape to get the flatfield correction
691 flatfield = H.reshape(ff_shape)
693 # Normalize the flatfield so its mean is centered around 1
694 flatfield_mean = np.mean(flatfield)
695 if flatfield_mean != 0:
696 flatfield /= flatfield_mean
698 print(
699 f"Estimated the flatfield ({flatfield.shape[0]}x{flatfield.shape[1]} pixels) for normalization with non-negative matrix factorization."
700 )
702 imageio.v3.imwrite(
703 shading_tiles_dir.parent.joinpath("nnmf_shading_correction.tif"),
704 flatfield.astype(np.float32),
705 )
707 # return flatfield
710import imageio.v3 as iio
711import tempfile
714def estimate_shading_correction_v2(
715 shading_tiles_dir: Path,
716 n_images_limit: int = 800,
717) -> np.ndarray:
718 """
719 Estimate the flatfield (shading correction) from a directory of TIFF tiles
720 using rank‐1 NMF, but without ever loading all images into RAM at once.
722 Parameters
723 ----------
724 shading_tiles_dir : Path
725 Directory containing the '.tif' shading‐tile images.
726 n_images_limit : int, optional
727 Maximum number of images to sample for the NMF fitting.
729 Returns
730 -------
731 flatfield : np.ndarray
732 2D array of the shading‐correction (mean scaled to 1).
733 """
735 # 1) Gather and possibly subsample the list of files
736 shading_images = ut.sorted_files(shading_tiles_dir, ".tif")
737 n_total = len(shading_images)
738 if n_total > n_images_limit:
739 print(
740 f"Found {n_total} images, sampling {n_images_limit} of them for memory‐savings."
741 )
742 shading_images = np.random.choice(
743 shading_images, n_images_limit, replace=False
744 )
745 else:
746 print(f"Found {n_total} images.")
748 n_images = len(shading_images)
750 # 2) Peek at the first image to get shape
751 first_img = iio.imread(shading_images[0]).astype(np.float32)
753 height, width, _ = first_img.shape
754 del first_img
756 # 3) Create a temporary file on disk and memory‐map it
757 tmp = tempfile.NamedTemporaryFile(
758 dir=shading_tiles_dir.parent, suffix=".dat", delete=False
759 )
760 tmp_filename = tmp.name
761 tmp.close()
763 # A memmap of shape (n_images, height*width) dtype float32
764 X = np.memmap(
765 tmp_filename,
766 dtype=np.float32,
767 mode="w+",
768 shape=(n_images, height * width),
769 )
771 # 4) Stream each TIFF into the memmap
772 for idx, img_path in enumerate(
773 tqdm(shading_images, desc="Building memmap")
774 ):
775 img = iio.imread(img_path).astype(np.float32)
776 X[idx, :] = img.ravel()
777 X.flush()
779 # 5) Fit a rank‐1 NMF
780 # This will only hold (n_images × 1) W and (1 × n_pixels) H in RAM
781 model = NMF(
782 n_components=1,
783 init="random",
784 random_state=0,
785 max_iter=200,
786 tol=1e-4,
787 )
788 W = model.fit_transform(X) # shape = (n_images, 1)
789 H = model.components_[0] # shape = (height*width,)
791 # 6) Tear down the memmap file
792 del X
793 os.remove(tmp_filename)
795 # 7) Reshape H back into a 2D flatfield and normalize its mean to 1
796 flatfield = H.reshape((height, width))
797 m = flatfield.mean()
798 if m != 0:
799 flatfield /= m
801 # 8) Optionally write it out
802 out_path = (
803 shading_tiles_dir.parent
804 / f"nnmf_shading_correction_{n_images_limit}.tif"
805 )
806 iio.imwrite(out_path, flatfield.astype(np.float32))
808 print(
809 f"Estimated flatfield of size {flatfield.shape} and wrote to {out_path}"
810 )
813# def save_shading_correction(shading_correction_img: np.ndarray, save_location: Path):
814# """Save the shading correction"""
816# # make the directory if it doesn't exist
817# parent_dir = save_location.parent.mkdir(parents=True,exist_ok=True)
819# # save
820# imageio.imwrite(save_location,shading_correction_img)
823@app.command(
824 cls=CustomCLICommand,
825 name="estimate_basic_shading_correction",
826 short_help="Estimate BaSiC (fiji plugin) flatfield",
827)
828def estimate_basic_shading_correction(
829 shading_tiles_dir: Path, fiji_app: Path
830) -> NDArray:
831 """To come"""
833 # write ImageJ macro
834 shading_macro = shading_tiles_dir.parent.joinpath(
835 f"{shading_tiles_dir.name}_shading_macro.ijm"
836 )
838 print(f"Estimating basic shading correction: {shading_macro}")
840 if platform == "win32":
841 shading_tiles_dir_str = str(shading_tiles_dir).replace(os.sep, "/")
842 else:
843 shading_tiles_dir_str = str(shading_tiles_dir)
845 with open(shading_macro, "w", encoding="utf-8") as f:
846 # f.write('setBatchMode(true);\n')
847 f.write(f'File.openSequence("{shading_tiles_dir_str}");\n')
848 # f.write("imageTitle = getTitle();\n")
849 f.write(
850 f'run("BaSiC ", "processing_stack={shading_tiles_dir.name} flat-field=None dark-field=None shading_estimation=[Estimate shading profiles] shading_model=[Estimate flat-field only (ignore dark-field)] setting_regularisationparametes=Manual temporal_drift=Ignore correction_options=[Compute shading only] lambda_flat=5 lambda_dark=0.50");\n'
851 )
852 # f.write('close(imageTitle);\n')
853 f.write(f'selectImage("Flat-field:{shading_tiles_dir.name}");\n')
854 f.write(
855 f'saveAs("Tiff", "{str(shading_tiles_dir.parent.joinpath(f"BaSiC_flatfield_estimate_{shading_tiles_dir.name}.tif")).replace(os.sep, "/")}");\n'
856 )
857 f.write('run("Close All");\n')
858 f.write('run("Quit");')
859 # f.write('setBatchMode(false);')
860 f.close()
862 # result = subprocess.run([str(fiji_app), "--run", str(shading_macro)], capture_output=True, text=True)
863 # print("STDOUT:", result.stdout)
864 # print("STDERR:", result.stderr)
866 ## Don't add "--headless", to not run in headless
867 # https://forum.image.sc/t/set-memory-from-command-line-upon-startup-headless/8668/4
868 run_macro = Popen(
869 [str(fiji_app), "-macro", str(shading_macro)],
870 shell=False,
871 stdout=PIPE,
872 stderr=PIPE,
873 )
875 # stream_data = run_macro.communicate()[0]
876 stdout, stderr = run_macro.communicate()
878 # print("STDOUT:", stdout.decode())
879 # print("STDERR:", stderr.decode())
881 if "error" in str(stdout).lower():
882 raise OSError(
883 f"Shading estimation failed on current slice, Fiji output:\n{stdout}"
884 )
887@app.command(
888 cls=CustomCLICommand,
889 name="reorganize_tiles_by_channel",
890 short_help="Reorganize tiles by channel keyword",
891)
892def reorganize_tiles_by_channel(
893 tiles_dir: Path, channel_keywords: List[str], move_files: bool = True
894):
895 """Reorganize by moving files based on keywords such as Pol, Bright, or Dark."""
897 dest_folder = tiles_dir.parent.expanduser()
899 for file_path in tiles_dir.rglob(
900 "*"
901 ): # Use rglob to find all files recursively
902 if file_path.is_file():
903 for keyword in channel_keywords:
904 if keyword in file_path.name:
905 keyword_dest_directory = dest_folder / keyword
906 keyword_dest_directory.mkdir(
907 parents=True, exist_ok=True
908 ) # Create directory if it doesn't exist
909 if move_files:
910 shutil.move(
911 file_path, keyword_dest_directory / file_path.name
912 )
913 else:
914 shutil.copy2(
915 file_path, keyword_dest_directory / file_path.name
916 )
917 break # Move to the next file after finding the first matching keyword
919 print("\tTiles reorganized by channels.\n")
921 # remove origin directory
922 if move_files:
923 ut.rmdir(tiles_dir)
926def combine_rgb_channels(
927 red_channel_path, green_channel_path, blue_channel_path, output_path
928):
929 """_summary_
931 Args:
932 red_channel_path (_type_): _description_
933 green_channel_path (_type_): _description_
934 blue_channel_path (_type_): _description_
935 output_path (_type_): _description_
936 """
937 # Read the individual channel images
938 r_channel = imageio.v3.imread(red_channel_path)
939 g_channel = imageio.v3.imread(green_channel_path)
940 b_channel = imageio.v3.imread(blue_channel_path)
942 # Check if all channels have the same shape
943 if (
944 r_channel.shape != g_channel.shape
945 or r_channel.shape != b_channel.shape
946 ):
947 print("Error: All channel images must have the same dimensions.")
948 return
950 # Stack the channels along the last dimension to create an RGB image
951 rgb_image = np.stack((r_channel, g_channel, b_channel), axis=-1)
953 # Save the combined image as a 32-bit RGB image
954 imageio.imwrite(
955 output_path, rgb_image, format="TIFF"
956 ) # Use TIFF for float support
959def separate_color_channels(input_dir: Path) -> list[Path]:
960 """To come"""
962 print(f"\tReading images in {input_dir}\n")
963 parent_dir = input_dir.parent
965 # Create directories for each channel in the parent directory
966 red_channel_dir = parent_dir / f"{input_dir.name}_channel_1_red"
967 green_channel_dir = parent_dir / f"{input_dir.name}_channel_2_green"
968 blue_channel_dir = parent_dir / f"{input_dir.name}_channel_3_blue"
970 red_channel_dir.mkdir(exist_ok=True)
971 green_channel_dir.mkdir(exist_ok=True)
972 blue_channel_dir.mkdir(exist_ok=True)
974 print(
975 f"Splitting RGB images into \n\t\t{red_channel_dir}\n\t\t{green_channel_dir}\n\t\t{blue_channel_dir}"
976 )
978 all_color_images_paths = ut.sorted_files(input_dir, ".tif")
980 for each_color_image_path in tqdm(all_color_images_paths):
981 # color_image = Image.open(each_color_image_path)
982 color_image = imageio.v3.imread(each_color_image_path)
984 # # Check if the image is in RGB mode
985 # if color_image.mode != 'RGB':
986 # print(f"The image is in {color_image.mode} mode, not RGB. Please provide an RGB image.")
987 # ut.rmdir(red_channel_dir)
988 # ut.rmdir(green_channel_dir)
989 # ut.rmdir(blue_channel_dir)
990 # return
992 # # Split the image into its RGB channels
993 # r, g, b = color_image.split()
995 # Check if the image is RGB
996 if color_image.ndim != 3 or color_image.shape[2] != 3:
997 print("The image is not in RGB format.")
998 ut.rmdir(red_channel_dir)
999 ut.rmdir(green_channel_dir)
1000 ut.rmdir(blue_channel_dir)
1001 return
1003 # Split the image into its RGB channels
1004 r_channel = color_image[:, :, 0] # Red channel
1005 g_channel = color_image[:, :, 1] # Green channel
1006 b_channel = color_image[:, :, 2] # Blue channel
1008 # Save each channel as an image
1009 imageio.imwrite(
1010 red_channel_dir / f"{each_color_image_path.stem}_channel1_red.tif",
1011 r_channel,
1012 )
1013 imageio.imwrite(
1014 green_channel_dir
1015 / f"{each_color_image_path.stem}_channel2_green.tif",
1016 g_channel,
1017 )
1018 imageio.imwrite(
1019 blue_channel_dir
1020 / f"{each_color_image_path.stem}_channel3_blue.tif",
1021 b_channel,
1022 )
1024 return [red_channel_dir, green_channel_dir, blue_channel_dir]