Coverage for src/precon3d/aligner.py: 26%
338 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 03:50 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 03:50 +0000
1# aligner.py is aa refactor of Alignment.py
3from sys import platform
4import os
5import time
6import shutil
7from pathlib import Path
8from enum import Enum
9from typing import NamedTuple
10import tempfile
11from subprocess import Popen, DEVNULL
12import numpy as np
13from PIL import Image, ImageEnhance
14import scipy.ndimage
15import scipy.optimize
16import scipy.optimize
18Image.MAX_IMAGE_PIXELS = None
19import imageio
20from skimage import feature, transform, registration, measure, exposure
21import matplotlib.pyplot as plt
22import typer
23from rich.progress import (
24 Progress,
25 SpinnerColumn,
26 BarColumn,
27 TextColumn,
28) # Import Rich progress componentsimport yaml # Assuming the configuration is in YAML format
30from pystackreg import StackReg
32from skimage.transform import pyramid_gaussian
34# from skimage.color import rgb2gray
36from precon3d.stitcher import FijiAttrs
37import precon3d.utility as ut
38import precon3d.factory as gen_types
40# pylint: disable=wildcard-import
41from precon3d.custom_types import *
42from precon3d._my_typer_cli import CustomCLIGroup, CustomCLICommand
44# CLI
45app = typer.Typer(
46 cls=CustomCLIGroup,
47 no_args_is_help=True,
48 short_help="Image registration/alignment methods",
49)
52class RegistrationMethod(Enum):
53 """
54 Sepcific registration type
56 Attributes supported
57 --------------------
58 TURBOREG: str
59 Fiji's turboreg method
60 SIFT: str
61 skimage SIFT (keypoints)
62 PHASECORR: str
63 skimage phasecorrelation (fft)
64 """
66 TURBOREG: str = "turboreg"
67 SIFT: str = "sift"
68 PHASECORR: str = "phasecorr"
71class AlignmentAttrs(NamedTuple):
72 """_summary_
74 Args:
75 NamedTuple (_type_): _description_
76 """
78 registration_method: RegistrationMethod = "turboreg"
79 use_subarea: bool = False
80 subarea_x: int = None
81 subarea_y: int = None
82 subarea_width: int = None
83 subarea_height: int = None
86def create_aligner_attrs(data: Dict[str, Any]) -> AlignmentAttrs:
87 """Create AlignmentAttrs from a dictionary."""
88 return AlignmentAttrs(
89 registration_method=str(data["aligner_attrs"]["registration_method"]),
90 use_subarea=bool(data["aligner_attrs"]["use_subarea"]),
91 subarea_x=int(data["aligner_attrs"]["subarea_x"]),
92 subarea_y=int(data["aligner_attrs"]["subarea_y"]),
93 subarea_width=int(data["aligner_attrs"]["subarea_width"]),
94 subarea_height=int(data["aligner_attrs"]["subarea_height"]),
95 )
98class DeltaTranslation(NamedTuple):
99 """Represents a translation in a 2D space.
101 This class is used to define a translation vector with respect to the x and y axes.
103 Args:
104 dx (float): The translation distance along the x-axis.
105 dy (float): The translation distance along the y-axis.
106 """
108 dx: float
109 dy: float
111 def __str__(self):
112 """Return a string representation of the DeltaTranslation instance."""
113 return f"(dx={self.dx}, dy={self.dy})"
116def apply_shift(
117 mov_img: np.ndarray, dt: DeltaTranslation, scale_factor: float = 1.0
118) -> np.array:
119 # Apply shift to the second image
120 # shift = np.array(dt.dx, dt.dy)
121 tform = transform.AffineTransform(
122 translation=[scale_factor * dt.dx, scale_factor * dt.dy]
123 )
124 registered_image = transform.warp(
125 mov_img, tform.inverse, clip=False, preserve_range=True
126 )
127 return registered_image
130def pad_arrays_to_match(ref_array: np.ndarray, mov_array: np.ndarray):
131 """Pad the smaller array among ref_array and mov_array to match the size of the larger array."""
132 # Calculate the size difference
133 delta_height = ref_array.shape[0] - mov_array.shape[0]
134 delta_width = ref_array.shape[1] - mov_array.shape[1]
136 # Determine padding for height and width
137 pad_height = abs(delta_height // 2), abs(
138 delta_height // 2 + delta_height % 2
139 )
140 pad_width = abs(delta_width // 2), abs(delta_width // 2 + delta_width % 2)
142 # Apply padding to the smaller array
143 if delta_height < 0 or delta_width < 0:
144 # If reference array is smaller, pad it
145 ref_array = np.pad(ref_array, (pad_height, pad_width), mode="constant")
146 else:
147 # If moving array is smaller or they are equal (no padding needed), pad moving array
148 mov_array = np.pad(mov_array, (pad_height, pad_width), mode="constant")
150 # Ensure both arrays have the same shape by padding the other side if necessary
151 if ref_array.shape != mov_array.shape:
152 # If there's still a difference, this means one dimension was equal and the other was not
153 # Determine new deltas
154 delta_height = ref_array.shape[0] - mov_array.shape[0]
155 delta_width = ref_array.shape[1] - mov_array.shape[1]
157 # New padding calculations
158 pad_height = abs(delta_height), 0
159 pad_width = abs(delta_width), 0
161 # Apply additional padding to the smaller array
162 if delta_height != 0:
163 if delta_height > 0:
164 # Moving array is smaller in height
165 mov_array = np.pad(
166 mov_array, (pad_height, (0, 0)), mode="constant"
167 )
168 else:
169 # Reference array is smaller in height
170 ref_array = np.pad(
171 ref_array, (pad_height, (0, 0)), mode="constant"
172 )
174 if delta_width != 0:
175 if delta_width > 0:
176 # Moving array is smaller in width
177 mov_array = np.pad(
178 mov_array, ((0, 0), pad_width), mode="constant"
179 )
180 else:
181 # Reference array is smaller in width
182 ref_array = np.pad(
183 ref_array, ((0, 0), pad_width), mode="constant"
184 )
186 return ref_array, mov_array
189from scipy.optimize import minimize
190import scipy
193def calc_shift_max(
194 ref_array, mov_array
195): # (reference_image: Path, moving_image: Path):
196 """Maximize the product of the images"""
198 # # Read and convert images to grayscale
199 # ref_img = Image.open(reference_image).convert("L")
200 # mov_img = Image.open(moving_image).convert("L")
202 # # Convert images to numpy arrays
203 # ref_array = np.array(ref_img)
204 # mov_array = np.array(mov_img)
206 # mov_array = scipy.ndimage.rotate(mov_array,15)
208 # # Pad
209 # ref_array, mov_array = pad_arrays_to_match(ref_image_gray, mov_image_gray)
211 # Define a function to calculate the negative product sum (since we'll be minimizing)
212 def neg_product_sum(shift, ref_array, mov_array):
213 """Calculate negative of product sum for given shift."""
214 dx, dy = int(shift[0]), int(shift[1])
215 shifted_mov_array = np.roll(mov_array, shift=(dx, dy), axis=(0, 1))
217 # For simplicity, we'll ignore the edges where the images don't overlap
218 # product_sum = np.sum(ref_array * shifted_mov_array)
219 # return -math.log(product_sum) # Negative for minimization
220 # Calculate the log of each product to avoid overflow, then sum
221 # Adding a small constant epsilon to avoid log(0)
222 epsilon = 1e-10
223 log_product_sum = np.sum(
224 np.log(np.abs(ref_array * shifted_mov_array) + epsilon)
225 )
227 return -log_product_sum # Return negative because we're minimizing
229 # Initial guess for the shift
230 initial_shift = [0, 0]
232 # Use scipy's minimize function to find the optimal shift
233 result = minimize(
234 neg_product_sum,
235 initial_shift,
236 args=(ref_array, mov_array),
237 method="Powell",
238 )
240 optimal_shift = result.x
241 print(f"Optimal shift: {optimal_shift}")
243 return optimal_shift
244 # breakpoint()
246 # # Apply the optimal shift to the moving image
247 # optimal_dx, optimal_dy = int(optimal_shift[0]), int(optimal_shift[1])
248 # optimal_shifted_mov_array = np.roll(mov_array, shift=(optimal_dx, optimal_dy), axis=(0, 1))
250 # # Convert back to PIL image for any further processing or saving
251 # optimal_shifted_mov_img = Image.fromarray(optimal_shifted_mov_array)
253 # return optimal_shifted_mov_img
256def pystackreg_shift(reference, moving):
257 """Optimize the shift between two images using pystackreg."""
259 # read in
260 ref_img = Image.open(reference).convert("L")
261 mov_img = Image.open(moving).convert("L")
263 # # Convert images to numpy arrays
264 ref_array = np.array(ref_img)
265 mov_array = np.array(mov_img)
267 sr = StackReg(StackReg.TRANSLATION)
269 # register images
270 tmats = sr.register(ref_array, mov_array)
272 return DeltaTranslation(dx=tmats[0, 2], dy=tmats[1, 2])
275def optimize_shift(ref_array, mov_array):
276 """Optimize the shift between two images using image pyramids."""
278 # Convert images to grayscale
279 # ref_image_gray = rgb2gray(ref_image)
280 # mov_image_gray = rgb2gray(mov_image)
282 reduction_factor = 10
284 # Generate image pyramids
285 pyramid_ref = tuple(
286 pyramid_gaussian(ref_array, downscale=reduction_factor)
287 )
288 pyramid_mov = tuple(
289 pyramid_gaussian(mov_array, downscale=reduction_factor)
290 )
292 # Start with an initial guess for the shift
293 initial_shift = np.array([0, 0])
295 # Iterate over pyramid levels starting from the smallest image
296 for ref, mov in zip(reversed(pyramid_ref), reversed(pyramid_mov)):
297 # Scale the initial shift for the current resolution
298 initial_shift *= reduction_factor
300 # Define the objective function for optimization
301 # def objective_function(shift):
302 # shifted_mov = np.roll(mov, shift.astype(int), axis=(0, 1))
303 # # Calculate the negative of the sum of squared differences (SSD)
304 # return -np.sum((ref - shifted_mov) ** 2)
305 # Define a function to calculate the negative product sum (since we'll be minimizing)
306 def objective_function(shift):
307 """Calculate negative of product sum for given shift."""
308 dx, dy = int(shift[0]), int(shift[1])
309 shifted_mov_array = np.roll(mov, shift=(dx, dy), axis=(0, 1))
311 # For simplicity, we'll ignore the edges where the images don't overlap
312 # product_sum = np.sum(ref_array * shifted_mov_array)
313 # return -math.log(product_sum) # Negative for minimization
314 # Calculate the log of each product to avoid overflow, then sum
315 # Adding a small constant epsilon to avoid log(0)
316 epsilon = 1e-10
317 log_product_sum = np.sum(
318 np.log(np.abs(ref * shifted_mov_array) + epsilon)
319 )
321 return -log_product_sum # Return negative because we're minimizing
323 # Optimize the shift
324 result = minimize(objective_function, initial_shift, method="Powell")
325 initial_shift = result.x
327 print(f"Optimal shift: {initial_shift}")
328 return initial_shift
331# import cupy as cp
334# def pyramid_down_gpu(image, downscale=10):
335# """Downsample an image by a given factor using CuPy."""
336# # Assuming 'image' is a CuPy array
337# # Simple downsampling, for demonstration purposes
338# return image[::downscale, ::downscale]
341# def generate_pyramid_gpu(image, levels):
342# """Generate an image pyramid as a list of images of decreasing sizes using CuPy."""
343# pyramid = [image]
344# for level in range(1, levels):
345# pyramid.append(pyramid_down_gpu(pyramid[level - 1]))
346# return pyramid
349# def optimize_shift_pyd_gpu(ref_image_gpu, mov_image_gpu, levels=3):
350# """Optimize the shift between two images using image pyramids and CuPy for GPU acceleration.
352# Find the optimal shift to align mov_image_gpu with ref_image_gpu.
354# Parameters:
355# - ref_image_gpu: CuPy array of the reference image.
356# - mov_image_gpu: CuPy array of the moving image.
358# Returns:
359# - Optimal shift as a tuple (y_shift, x_shift).
360# """
361# # Convert images to grayscale and to CuPy arrays
362# # ref_image_gpu = cp.asarray(ref_image)
363# # mov_image_gpu = cp.asarray(mov_image)
365# # Generate image pyramids
366# pyramid_ref = generate_pyramid_gpu(ref_image_gpu, levels)
367# pyramid_mov = generate_pyramid_gpu(mov_image_gpu, levels)
369# initial_shift = cp.array([0, 0], dtype=cp.float64)
371# for level in reversed(range(levels)):
372# # Scale the initial shift for the current resolution
373# initial_shift *= 10
374# ref = pyramid_ref[level]
375# mov = pyramid_mov[level]
377# # Define the objective function for optimization
378# def objective_function(shift):
379# # Ensure shift is integer for indexing
380# shift = cp.asarray(shift).astype(cp.int32)
381# shifted_mov = cp.roll(mov, shift, axis=(0, 1))
382# # Calculate the negative of the sum of squared differences (SSD)
383# # ssd = -cp.sum((ref - shifted_mov) ** 2).get() # Use .get() to transfer result to CPU
384# # return ssd
385# # Adding a small constant epsilon to avoid log(0)
386# epsilon = 1e-5
387# log_product_sum = cp.sum(
388# cp.log(cp.abs(ref * shifted_mov) + epsilon)
389# ).get()
391# return -log_product_sum # Return negative because we're minimizing
393# result = scipy.optimize.basinhopping(
394# objective_function,
395# initial_shift.get(),
396# niter=200,
397# stepsize=50,
398# niter_success=50,
399# )
401# # Optimize the shift using scipy (CPU-based)
402# # Note: The objective function transfers the final SSD value back to the CPU
403# # result = minimize(objective_function, initial_shift.get(), method='Powell') # Use .get() for CPU compatibility
404# initial_shift = cp.array(
405# result.x
406# ) # Transfer optimized shift back to GPU
408# return initial_shift.get() # Return final shift as a NumPy array
411# def optimize_shift_gpu(ref_image_gpu, mov_image_gpu, levels=3):
412# """Optimize the shift between two images using image pyramids and CuPy for GPU acceleration.
414# Find the optimal shift to align mov_image_gpu with ref_image_gpu.
416# Parameters:
417# - ref_image_gpu: CuPy array of the reference image.
418# - mov_image_gpu: CuPy array of the moving image.
420# Returns:
421# - Optimal shift as a tuple (y_shift, x_shift).
422# """
423# # Convert images to grayscale and to CuPy arrays
424# # ref_image_gpu = cp.asarray(ref_image)
425# # mov_image_gpu = cp.asarray(mov_image)
427# # Generate image pyramids
428# # pyramid_ref = generate_pyramid_gpu(ref_image_gpu, levels)
429# # pyramid_mov = generate_pyramid_gpu(mov_image_gpu, levels)
431# ref = ref_image_gpu
432# mov = mov_image_gpu
434# initial_shift = cp.array([0, 0], dtype=cp.float64)
436# # Define the objective function for optimization
437# def objective_function(shift):
438# # Ensure shift is integer for indexing
439# shift = cp.asarray(shift).astype(cp.int32)
440# shifted_mov = cp.roll(mov, shift, axis=(0, 1))
441# # Calculate the negative of the sum of squared differences (SSD)
442# # ssd = -cp.sum((ref - shifted_mov) ** 2).get() # Use .get() to transfer result to CPU
443# # return ssd
444# # Adding a small constant epsilon to avoid log(0)
445# epsilon = 1e-5
446# log_product_sum = cp.sum(
447# cp.log(cp.abs(ref * shifted_mov) + epsilon)
448# ).get()
450# return -log_product_sum # Return negative because we're minimizing
452# result = scipy.optimize.basinhopping(
453# objective_function,
454# initial_shift.get(),
455# niter=200,
456# stepsize=50,
457# niter_success=20,
458# )
460# return result.x
461# # Optimize the shift using scipy (CPU-based)
462# # Note: The objective function transfers the final SSD value back to the CPU
463# # result = minimize(objective_function, initial_shift.get(), method='Nelder-Mead') # Use .get() for CPU compatibility
464# # initial_shift = cp.array(result.x) # Transfer optimized shift back to GPU
466# result = minimize(
467# objective_function, initial_shift.get(), method="Nelder-Mead"
468# ) # method="BFGS")#method='Powell') # Use .get() for CPU compatibility
470# brute_range = ((-100, 100), (-100, 100))
471# result = scipy.optimize.brute(func=objective_function, ranges=brute_range)
472# return result
474# # initial_shift = cp.array(result.x) # Transfer optimized shift back to GPU
476# # return initial_shift.get() # Return final shift as a NumPy array
479def calc_shifts_phase_corr(
480 reference_image: Path,
481 moving_image: Path,
482 alignment_attrs: AlignmentAttrs,
483) -> DeltaTranslation:
484 """_summary_
486 Args:
487 reference_image (Path): _description_
488 moving_image (Path): _description_
489 alignment_attrs (AlignmentAttrs): _description_
491 Returns:
492 DeltaTranslation: _description_
493 """
495 # Read and convert images to grayscale
496 ref_img = Image.open(reference_image).convert("L")
497 mov_img = Image.open(moving_image).convert("L")
499 # Crop the images if use_subarea is set
500 if alignment_attrs.use_subarea:
501 # Define the cropping box
502 crop_box = (
503 alignment_attrs.subarea_x,
504 alignment_attrs.subarea_y,
505 alignment_attrs.subarea_x + alignment_attrs.subarea_width,
506 alignment_attrs.subarea_y + alignment_attrs.subarea_height,
507 )
508 ref_img = ref_img.crop(crop_box)
509 mov_img = mov_img.crop(crop_box)
511 ref_img = np.array(ref_img)
512 mov_img = np.array(mov_img)
514 shift, _, _ = registration.phase_cross_correlation(ref_img, mov_img)
516 delta_translation = DeltaTranslation(dx=shift[1], dy=shift[0])
518 return delta_translation
521def coarse_calc_shifts_phase_corr(
522 reference_image: Path,
523 moving_image: Path,
524 alignment_attrs: AlignmentAttrs,
525 scale_factor: float = 0.05,
526) -> DeltaTranslation:
527 """_summary_
529 Args:
530 reference_image (Path): _description_
531 moving_image (Path): _description_
532 alignment_attrs (AlignmentAttrs): _description_
534 Returns:
535 DeltaTranslation: _description_
536 """
538 # Read and convert images to grayscale
539 ref_img = Image.open(reference_image).convert("L")
540 mov_img = Image.open(moving_image).convert("L")
542 # Crop the images if use_subarea is set
543 if alignment_attrs.use_subarea:
544 # Define the cropping box
545 crop_box = (
546 alignment_attrs.subarea_x,
547 alignment_attrs.subarea_y,
548 alignment_attrs.subarea_x + alignment_attrs.subarea_width,
549 alignment_attrs.subarea_y + alignment_attrs.subarea_height,
550 )
551 ref_img = ref_img.crop(crop_box)
552 mov_img = mov_img.crop(crop_box)
554 ref_img_scaled = scipy.ndimage.zoom(
555 np.array(ref_img), (scale_factor, scale_factor), order=3
556 )
557 mov_img_scaled = scipy.ndimage.zoom(
558 np.array(mov_img), (scale_factor, scale_factor), order=3
559 )
561 ref_img_8bit = ut.convert_gray_to_8bit(ref_img_scaled)
562 mov_img_8bit = ut.convert_gray_to_8bit(mov_img_scaled)
564 shift, _, _ = registration.phase_cross_correlation(
565 ref_img_8bit, mov_img_8bit
566 )
568 delta_translation = DeltaTranslation(
569 dx=shift[1] / scale_factor, dy=shift[0] / scale_factor
570 )
572 print(f"{moving_image.name} shifts {delta_translation}")
574 return delta_translation
577def calc_shifts_phase_corr_pyramid(
578 reference_image: Path,
579 moving_image: Path,
580 alignment_attrs: AlignmentAttrs,
581 max_layer=3,
582) -> DeltaTranslation:
583 """_summary_
585 Args:
586 reference_image (Path): _description_
587 moving_image (Path): _description_
588 alignment_attrs (AlignmentAttrs): _description_
590 Returns:
591 DeltaTranslation: _description_
592 """
594 # Read and convert images to grayscale
595 ref_img = Image.open(reference_image).convert("L")
596 mov_img = Image.open(moving_image).convert("L")
598 # Crop the images if use_subarea is set
599 if alignment_attrs.use_subarea:
600 # Define the cropping box
601 crop_box = (
602 alignment_attrs.subarea_x,
603 alignment_attrs.subarea_y,
604 alignment_attrs.subarea_x + alignment_attrs.subarea_width,
605 alignment_attrs.subarea_y + alignment_attrs.subarea_height,
606 )
607 ref_img = ref_img.crop(crop_box)
608 mov_img = mov_img.crop(crop_box)
610 ref_img = np.array(ref_img)
611 mov_img = np.array(mov_img)
613 # Generate Gaussian pyramids for both images
614 pyramid1 = tuple(transform.pyramid_gaussian(ref_img, max_layer=max_layer))
615 pyramid2 = tuple(transform.pyramid_gaussian(mov_img, max_layer=max_layer))
617 shift = np.array([0, 0])
618 for layer in range(max_layer, -1, -1):
619 # Calculate shift at current pyramid level
620 shift_temp, _, _ = registration.phase_cross_correlation(
621 pyramid1[layer], pyramid2[layer], upsample_factor=10
622 )
623 # Upscale shift for next level
624 shift = (shift + shift_temp) * 2 if layer > 0 else shift + shift_temp
626 # average
627 shift = shift / (max_layer + 1)
629 delta_translation = DeltaTranslation(dx=shift[1], dy=shift[0])
631 return delta_translation
634def calc_shifts_using_keypoints(
635 reference_image: Path,
636 moving_image: Path,
637 alignment_attrs: AlignmentAttrs,
638 show_plots: bool = False,
639) -> DeltaTranslation:
640 """_summary_
642 Args:
643 reference_image (Path): _description_
644 moving_image (Path): _description_
645 alignment_attrs (AlignmentAttrs): _description_
646 fiji_attrs (FijiAttrs): _description_
648 Returns:
649 DeltaTranslation: _description_
650 """
652 # Read and convert images to grayscale
653 ref_img = Image.open(reference_image) # .convert("L")
654 mov_img = Image.open(moving_image) # .convert("L")
656 # Crop the images if use_subarea is set
657 if alignment_attrs.use_subarea:
658 # Define the cropping box
659 crop_box = (
660 alignment_attrs.subarea_x,
661 alignment_attrs.subarea_y,
662 alignment_attrs.subarea_x + alignment_attrs.subarea_width,
663 alignment_attrs.subarea_y + alignment_attrs.subarea_height,
664 )
665 ref_img = ref_img.crop(crop_box)
666 mov_img = mov_img.crop(crop_box)
668 # convert to gray (8bit)
669 ref_img = Image.fromarray(
670 ut.convert_gray_to_8bit(np.array(ref_img), gamma=0.5)
671 ).convert("L")
672 mov_img = Image.fromarray(
673 ut.convert_gray_to_8bit(np.array(mov_img), gamma=0.5)
674 ).convert("L")
676 # Create a Contrast enhancer object
677 enhancer_ref_img = ImageEnhance.Contrast(ref_img)
678 enhancer_mov_img = ImageEnhance.Contrast(mov_img)
680 # Enhance the contrast (e.g., by a factor of 3)
681 ref_img = enhancer_ref_img.enhance(3)
682 mov_img = enhancer_mov_img.enhance(3)
684 ref_img = np.array(ref_img)
685 mov_img = np.array(mov_img)
687 # Enhance contrast using CLAHE
688 # ref_img = exposure.equalize_adapthist(ref_img, clip_limit=0.03)
689 # mov_img = exposure.equalize_adapthist(mov_img, clip_limit=0.03)
691 # Detect SIFT keypoints and descriptors
692 sift = feature.SIFT()
693 sift.detect_and_extract(ref_img)
694 keypoints1 = sift.keypoints
695 descriptors1 = sift.descriptors
697 sift.detect_and_extract(mov_img)
698 keypoints2 = sift.keypoints
699 descriptors2 = sift.descriptors
701 # # Detect ORB keypoints and descriptors
702 # orb = feature.ORB(n_keypoints=500, fast_threshold=0.05)
704 # orb.detect_and_extract(ref_img)
705 # keypoints1 = orb.keypoints
706 # descriptors1 = orb.descriptors
708 # orb.detect_and_extract(mov_img)
709 # keypoints2 = orb.keypoints
710 # descriptors2 = orb.descriptors
712 # Match descriptors using Hamming distance
713 matches = feature.match_descriptors(
714 descriptors1, descriptors2, cross_check=True
715 )
717 src = keypoints2[matches[:, 1]][:, ::-1]
718 dst = keypoints1[matches[:, 0]][:, ::-1]
720 # Use RANSAC to estimate a robust transform
721 model_robust, inliers = measure.ransac(
722 (src, dst),
723 transform.EuclideanTransform,
724 min_samples=4,
725 residual_threshold=2,
726 max_trials=1000,
727 )
729 # Extract matched keypoints
730 # matched_keypoints1 = keypoints1[matches[:, 0]]
731 # matched_keypoints2 = keypoints2[matches[:, 1]]
733 # # Estimate translation using RANSAC
734 # model_robust = transform.estimate_transform('euclidean', matched_keypoints1, matched_keypoints2)
735 # # model_robust, _ = transform.estimate_transform('euclidean', matched_keypoints1, matched_keypoints2, return_inliers=True)
737 # # Extract translation amount
738 # translation = model_robust.translation
740 best_translation = DeltaTranslation(
741 model_robust.translation[0], model_robust.translation[1]
742 )
744 # Visualize matches
745 if show_plots:
747 print(f"Estimated translation: {best_translation}")
748 inlier_matches = matches[inliers]
750 fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
752 # Plot the images next to each other
753 img3 = np.concatenate((ref_img, mov_img), axis=1)
754 ax.imshow(img3, cmap="gray")
755 ax.plot(
756 keypoints1[matches[:, 0], 1],
757 keypoints1[matches[:, 0], 0],
758 "ro",
759 markersize=4,
760 alpha=0.1,
761 )
762 ax.plot(
763 keypoints2[matches[:, 1], 1] + ref_img.shape[1],
764 keypoints2[matches[:, 1], 0],
765 "ro",
766 markersize=4,
767 alpha=0.1,
768 )
770 ax.plot(
771 keypoints1[inlier_matches[:, 0], 1],
772 keypoints1[inlier_matches[:, 0], 0],
773 "co",
774 markersize=8,
775 alpha=0.2,
776 )
777 ax.plot(
778 keypoints2[inlier_matches[:, 1], 1] + ref_img.shape[1],
779 keypoints2[inlier_matches[:, 1], 0],
780 "mo",
781 markersize=8,
782 alpha=0.2,
783 )
785 # Draw lines between matched keypoints
786 for match in inlier_matches:
787 pt1 = (keypoints1[match[0]][1], keypoints1[match[0]][0])
788 pt2 = (
789 keypoints2[match[1]][1] + ref_img.shape[1],
790 keypoints2[match[1]][0],
791 )
792 ax.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]], "g", alpha=0.1)
793 # ax.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]], "go", alpha=0.05)
795 # Turn off axis
796 ax.axis("off")
797 plt.title(f"Estimated translation: {best_translation}")
798 plt.tight_layout()
799 plt.show()
801 # breakpoint()
803 return best_translation
806# def calc_shifts_turboreg(
807# reference_image: Path,
808# moving_image: Path,
809# alignment_attrs: AlignmentAttrs,
810# fiji_attrs: FijiAttrs,
811# ) -> DeltaTranslation:
812# """
813# Calculate shifts using TurboReg.
815# Parameters
816# ----------
817# reference_image : Path
818# Path to the reference image.
819# moving_image : Path
820# Path to the moving image.
821# alignment_attrs : AlignmentAttrs
822# Attributes related to alignment, including cropping.
823# fiji_attrs : FijiAttrs
824# Attributes related to Fiji application settings.
826# Returns
827# -------
828# DeltaTranslation
829# The calculated shifts (dx, dy).
830# """
832# with tempfile.TemporaryDirectory() as temp_dir:
833# # Define the destination path for the temporary image
834# temp_reference_image = os.path.join(temp_dir, os.path.basename(reference_image))
835# temp_moving_image = os.path.join(temp_dir, os.path.basename(moving_image))
837# # Read the image using imageio
838# ref_img = imageio.v3.imread(reference_image)
839# mov_img = imageio.v3.imread(moving_image)
841# # Convert to gray if color
842# if ref_img.ndim == 2 and mov_img.ndim == 2:
843# pass
844# elif ref_img.ndim == 3 and ref_img.shape[-1] == 3 and mov_img.ndim == 3 and mov_img.shape[-1] == 3:
845# ref_img = skimage.color.rgb2gray(ref_img)
846# mov_img = skimage.color.rgb2gray(mov_img)
847# elif ref_img.ndim == 3 and ref_img.shape[-1] == 1 and mov_img.ndim == 3 and mov_img.shape[-1] == 1:
848# ref_img = ref_img.squeeze()
849# mov_img = mov_img.squeeze()
850# else:
851# raise TypeError(f'Images must be 3-channel color or gray, not shaped as {ref_img.shape} and {mov_img.shape}')
854def calc_shifts_turboreg(
855 reference_image: Path,
856 moving_image: Path,
857 alignment_attrs: AlignmentAttrs,
858 fiji_attrs: FijiAttrs,
859) -> DeltaTranslation:
860 """
861 Calculate shifts using TurboReg.
863 Parameters
864 ----------
865 reference_image : Path
866 Path to the reference image.
867 moving_image : Path
868 Path to the moving image.
869 alignment_attrs : AlignmentAttrs
870 Attributes related to alignment, including cropping.
871 fiji_attrs : FijiAttrs
872 Attributes related to Fiji application settings.
874 Returns
875 -------
876 DeltaTranslation
877 The calculated shifts (dx, dy).
878 """
880 with tempfile.TemporaryDirectory() as temp_dir:
882 # Define the destination path for the temporary image
883 temp_reference_image = os.path.join(
884 temp_dir, os.path.basename(reference_image)
885 )
886 temp_moving_image = os.path.join(
887 temp_dir, os.path.basename(moving_image)
888 )
890 # Read and convert images to grayscale
891 ref_img = Image.open(reference_image) # .convert("L")
892 mov_img = Image.open(moving_image) # .convert("L")
894 # Crop the images if use_subarea is set
895 if alignment_attrs.use_subarea:
896 # Define the cropping box
897 crop_box = (
898 alignment_attrs.subarea_x,
899 alignment_attrs.subarea_y,
900 alignment_attrs.subarea_x + alignment_attrs.subarea_width,
901 alignment_attrs.subarea_y + alignment_attrs.subarea_height,
902 )
903 ref_img = ref_img.crop(crop_box)
904 mov_img = mov_img.crop(crop_box)
906 # adjust the contrast
908 # Create a Contrast enhancer object
909 enhancer_ref_img = ImageEnhance.Contrast(ref_img)
910 enhancer_mov_img = ImageEnhance.Contrast(mov_img)
912 # Enhance the contrast (e.g., by a factor of 3)
913 ref_img = enhancer_ref_img.enhance(3)
914 mov_img = enhancer_mov_img.enhance(3)
916 # Save the processed images to the temporary directory
917 ref_img.save(temp_reference_image)
918 mov_img.save(temp_moving_image)
920 # Get the width and height of the processed image
921 height, width = ref_img.size
923 # # Read the image using Pillow
924 # ref_img = Image.open(reference_image)
925 # mov_img = Image.open(moving_image)
927 # # Convert to grayscale if the images are in color
928 # if ref_img.mode == 'L' and mov_img.mode == 'L':
929 # pass # Both images are already grayscale
930 # elif ref_img.mode in ['RGB', 'RGBA'] and mov_img.mode in ['RGB', 'RGBA']:
931 # ref_img = ref_img.convert('L') # Convert to grayscale
932 # mov_img = mov_img.convert('L') # Convert to grayscale
933 # else:
934 # raise TypeError(f'Images must be 3-channel color or gray, not shaped as {ref_img.size} and {mov_img.size}')
936 # # Save the temporary grayscale images if needed
937 # ref_img.save(temp_reference_image)
938 # mov_img.save(temp_moving_image)
940 # if alignment_attrs.use_subarea:
941 # # Crop the image using NumPy slicing
942 # cropped_ref_img = ref_img[
943 # alignment_attrs.subarea_y : alignment_attrs.subarea_y
944 # + alignment_attrs.subarea_height,
945 # alignment_attrs.subarea_x : alignment_attrs.subarea_x
946 # + alignment_attrs.subarea_width,
947 # ]
949 # cropped_mov_img = mov_img[
950 # alignment_attrs.subarea_y : alignment_attrs.subarea_y
951 # + alignment_attrs.subarea_height,
952 # alignment_attrs.subarea_x : alignment_attrs.subarea_x
953 # + alignment_attrs.subarea_width,
954 # ]
956 # # save the cropped image to the temporary directory
957 # imageio.v3.imwrite(temp_reference_image, cropped_ref_img)
958 # imageio.v3.imwrite(temp_moving_image, cropped_mov_img)
960 # else:
961 # # save the uncropped image to the temporary directory
962 # imageio.v3.imwrite(temp_reference_image, ref_img)
963 # imageio.v3.imwrite(temp_moving_image, mov_img)
965 # # Get the width and height of the cropped image
966 # height = imageio.v3.imread(temp_reference_image).shape[0]
967 # width = imageio.v3.imread(temp_reference_image).shape[1]
969 # define the location of the shift file
970 temp_calculated_shifts = os.path.join(
971 temp_dir, "calculated_shifts.csv"
972 )
974 # modify file path for the macro if using windows
975 if platform == "win32":
976 temp_reference_image = temp_reference_image.replace(os.sep, "/")
977 temp_moving_image = temp_moving_image.replace(os.sep, "/")
978 temp_calculated_shifts = temp_calculated_shifts.replace(
979 os.sep, "/"
980 )
982 # write ImageJ macro
983 turbo_reg_macro = os.path.join(temp_dir, "turbo_reg_macro.ijm")
984 Path(turbo_reg_macro).touch() # Creates an empty file
985 with open(turbo_reg_macro, "w", encoding="utf-8") as f:
986 f.write(f'target = "{temp_reference_image}"; //reference\n')
987 f.write(f'source = "{temp_moving_image}"; //thing to be aligned\n')
988 f.write(f"width = {width};\n")
989 f.write(f"height = {height};\n")
991 f.write('run("TurboReg ","-align " +\n')
992 f.write('\t\t\t\t"-file " + source + \n')
993 f.write(
994 '\t\t\t\t" 0 0 " + width + " " + height + //cropping (start x, start y, width, height)\n'
995 )
996 f.write('\t\t\t\t" -file " + target + \n')
997 f.write(
998 '\t\t\t\t" 0 0 " + width + " " + height + //cropping (start x, start y, width, height)\n'
999 )
1000 f.write('\t\t\t\t" -translation " +\n')
1001 f.write(
1002 '\t\t\t\twidth/2 + " " + height/2 + " " + //landmark x,y (center of image)\n'
1003 )
1004 f.write(
1005 '\t\t\t\twidth/2 + " " + height/2 + //landmark x,y (center of image)\n'
1006 )
1007 f.write('\t\t\t\t" -hideOutput");\n')
1008 # f.write('\t\t\t\t" -showOutput");\n')
1010 # interpret result
1011 f.write('sourceX0 = getResult("sourceX", 0);\n')
1012 f.write('sourceY0 = getResult("sourceY", 0);\n')
1013 f.write('targetX0 = getResult("targetX", 0);\n')
1014 f.write('targetY0 = getResult("targetY", 0);\n')
1015 f.write("dx = targetX0 - sourceX0;\n")
1016 f.write("dy = targetY0 - sourceY0;\n")
1018 # write out teh result
1019 f.write(f'f = File.open("{temp_calculated_shifts}");\n')
1020 f.write('print(f, d2s(dx,2) + ", " + d2s(dy,3));\n')
1021 f.close()
1023 fiji_app = fiji_attrs.fiji_app
1024 ## managing memory in headless
1025 # https://forum.image.sc/t/set-memory-from-command-line-upon-startup-headless/8668/4
1026 run_macro = Popen(
1027 [fiji_app, "--headless", "-macro", turbo_reg_macro],
1028 shell=False,
1029 stdout=DEVNULL,
1030 stderr=DEVNULL,
1031 )
1032 stream_data = run_macro.communicate()[0]
1033 if "error" in str(stream_data).lower():
1034 raise OSError(
1035 f"Stitching failed on current slice, Fiji output:\n{stream_data}"
1036 )
1038 ##convert to numpy file
1039 shift_data = np.genfromtxt(
1040 temp_calculated_shifts, delimiter=","
1041 ) # [:, 2:].astype(int)
1042 delta_translation = DeltaTranslation(
1043 dx=shift_data[0], dy=shift_data[1]
1044 )
1046 return delta_translation
1049# def apply_all_shifts_deprecated(image_dir: Path, shifts_path: Path):
1050# """
1051# to come
1052# """
1054# image_file_list = ut.sorted_files(image_dir, 'tif')
1056# calculated_shifts = np.load(shifts_path)
1057# c_shifts = np.cumsum(calculated_shifts,axis=0)
1059# # check sizes make sense
1060# if len(image_file_list)-1 != len(calculated_shifts):
1061# raise ValueError(f"Found {len(image_file_list)} files, expected {len(calculated_shifts)} from calculated shifts")
1064# save_dir = image_dir.parent.joinpath(f"aligned_xy_{image_dir.stem}")
1065# save_dir.mkdir(exist_ok=True)
1067# # copy first image
1068# shutil.copyfile(image_file_list[0],save_dir.joinpath(image_file_list[0].name))
1070# for i, each_shift in enumerate(c_shifts, start=1):
1071# # convert to a DeltaTranslation
1072# dt = DeltaTranslation(dx=each_shift[0],dy=each_shift[1])
1073# orig_img = np.array(Image.open(image_file_list[i]))
1074# shifted_orig_img = apply_shift(orig_img,dt)
1075# new_img = Image.fromarray(shifted_orig_img.astype(orig_img.dtype))
1076# new_img.save(save_dir.joinpath(image_file_list[i].name))
1079@app.command(
1080 cls=CustomCLICommand,
1081 name="apply_all_shifts",
1082 short_help="apply known shifts",
1083)
1084def apply_all_shifts(
1085 image_dir: Path, shifts_path: Path, scale_factor: float = 1.0
1086):
1087 """
1088 to come
1089 """
1091 # cast to Path object if not already a Path object
1092 if not isinstance(image_dir, Path):
1093 image_dir = Path(image_dir)
1094 if not isinstance(shifts_path, Path):
1095 shifts_path = Path(shifts_path)
1097 image_file_list = ut.sorted_files(image_dir, "tif")
1099 calculated_shifts = np.load(shifts_path)
1100 c_shifts = np.cumsum(calculated_shifts, axis=0)
1102 # check sizes make sense
1103 # if len(image_file_list)-1 != len(calculated_shifts):
1104 # raise ValueError(f"Found {len(image_file_list)} files, expected {len(calculated_shifts)} from calculated shifts")
1105 if len(image_file_list) % (len(calculated_shifts) + 1) != 0:
1106 raise ValueError(
1107 f"Oops, found {len(image_file_list)} files, and {len(calculated_shifts)} from calculated shifts"
1108 )
1109 else:
1110 print(
1111 f"Found {len(image_file_list)} files, and {len(calculated_shifts)} from calculated shifts"
1112 )
1114 save_dir = image_dir.parent.joinpath(f"aligned_xy_{image_dir.stem}")
1115 save_dir.mkdir(exist_ok=True)
1117 # copy first image
1118 # first_imgs = image_file_list[0::13]
1119 # for each_fist_img in first_imgs:
1120 # shutil.copyfile(each_fist_img,save_dir.joinpath(each_fist_img.name))
1122 # Use list comprehension to exclude every 1st, 14th, 27th, etc., element
1123 # filtered_list = [item for index, item in enumerate(image_file_list, start=1) if (index - 1) % 13 != 0]
1125 shutil.copyfile(
1126 image_file_list[0], save_dir.joinpath(image_file_list[0].name)
1127 )
1129 for shift_idx, each_img in enumerate(image_file_list[1:]):
1131 shift = c_shifts[shift_idx]
1133 # convert to a DeltaTranslation
1134 dt = DeltaTranslation(dx=shift[0], dy=shift[1])
1135 orig_img = np.array(Image.open(each_img))
1136 shifted_orig_img = apply_shift(orig_img, dt, scale_factor)
1137 new_img = Image.fromarray(shifted_orig_img.astype(orig_img.dtype))
1138 new_img.save(save_dir.joinpath(each_img.name))
1140 print(f"\tapplied shift to {each_img.name}")
1143def list_to_numpy_array(delta_translations):
1144 # Extract dx and dy from each DeltaTranslation object and create a list of tuples
1145 data = [(dt.dx, dt.dy) for dt in delta_translations]
1146 # Convert the list of tuples into a NumPy array
1147 return np.array(data)
1150def numpy_array_to_list(numpy_array):
1151 # Create a list of DeltaTranslation objects from the NumPy array
1152 return [DeltaTranslation(dx=row[0], dy=row[1]) for row in numpy_array]
1155# def run_aligner(configfile: Path):
1156# """align image stacks using the user defined yaml
1158# Args:
1159# configfile (Path): _description_
1160# """
1162# # cast to Path object if not already a Path object
1163# if not isinstance(configfile, Path):
1164# configfile = Path(configfile)
1166# # read yaml content
1167# config_dict = ut.read_config(configfile)
1169# user_aligner_attrs = create_aligner_attrs(config_dict)
1170# general_attrs = precon3d.factory.create_general_attrs(config_dict)
1171# ref_channel = str(config_dict["ref_channel"])
1173# # check the ref channel is a folder of the general_attrs:input_directory
1176# def align_scene_using_ref(config: Dict):
1177# """
1178# to come
1179# """
1182def align_image_stack_turboreg(
1183 image_dir: Path, alignment_attrs: AlignmentAttrs, fiji_attrs: FijiAttrs
1184):
1185 """_summary_"""
1187 # align
1188 if image_dir.is_dir():
1190 processed_imagelist = ut.sorted_files(image_dir, ".tif")
1192 print(f"Aligning {len(processed_imagelist)} images in {image_dir}")
1194 optimal_shifts = [
1195 calc_shifts_turboreg(
1196 ref_image, mov_image, alignment_attrs, fiji_attrs
1197 )
1198 for ref_image, mov_image in zip(
1199 processed_imagelist, processed_imagelist[1:]
1200 )
1201 ]
1203 calculated_shifts = image_dir.parent.joinpath(
1204 f"{image_dir.name}_optimized_shifts_turboreg.npy"
1205 )
1206 np.save(calculated_shifts, list_to_numpy_array(optimal_shifts))
1208 # apply shifts
1209 apply_all_shifts(image_dir, calculated_shifts)
1211 print("Alignment done.")
1214def align_image_stack_sift(image_dir: Path, alignment_attrs: AlignmentAttrs):
1215 """_summary_"""
1217 # align
1218 if image_dir.is_dir():
1220 processed_imagelist = ut.sorted_files(image_dir, ".tif")
1222 print(f"Aligning {len(processed_imagelist)} images in {image_dir}")
1224 optimal_shifts = [
1225 calc_shifts_using_keypoints(ref_image, mov_image, alignment_attrs)
1226 for ref_image, mov_image in zip(
1227 processed_imagelist, processed_imagelist[1:]
1228 )
1229 ]
1231 calculated_shifts = image_dir.parent.joinpath(
1232 f"{image_dir.name}_optimized_shifts_sift.npy"
1233 )
1234 np.save(calculated_shifts, list_to_numpy_array(optimal_shifts))
1236 # apply shifts
1237 apply_all_shifts(image_dir, calculated_shifts)
1239 print("Alignment done.")
1242def align_image_stack_phasecorr(
1243 image_dir: Path, alignment_attrs: AlignmentAttrs
1244):
1245 """_summary_"""
1247 # align
1248 if image_dir.is_dir():
1250 processed_imagelist = ut.sorted_files(image_dir, ".tif")
1252 print(f"Aligning {len(processed_imagelist)} images in {image_dir}")
1254 optimal_shifts = [
1255 calc_shifts_phase_corr(ref_image, mov_image, alignment_attrs)
1256 for ref_image, mov_image in zip(
1257 processed_imagelist, processed_imagelist[1:]
1258 )
1259 ]
1261 calculated_shifts = image_dir.parent.joinpath(
1262 f"{image_dir.name}_optimized_shifts_phasecorr.npy"
1263 )
1264 np.save(calculated_shifts, list_to_numpy_array(optimal_shifts))
1266 # apply shifts
1267 apply_all_shifts(image_dir, calculated_shifts)
1269 print("Alignment done.")
1272@app.command(
1273 cls=CustomCLICommand,
1274 name="process_images",
1275 short_help="align images from user defined config",
1276)
1277def process_images(config_filepath: Path):
1278 """
1279 align images from user defined config
1280 """
1282 # Read the dict in the config file
1283 config_dict = ut.read_config(config_filepath)
1284 alignment_attrs = create_aligner_attrs(config_dict)
1285 general_attrs = gen_types.create_general_attrs(config_dict)
1287 if not ut.valid_enum_entry(
1288 alignment_attrs.registration_method, RegistrationMethod
1289 ):
1290 raise KeyError(
1291 f"Unsupported registration method {alignment_attrs.registration_method}, supported types include {[i.value for i in RegistrationMethod]}"
1292 )
1293 print(f"Using: {alignment_attrs.registration_method} method to align")
1295 tic = time.perf_counter()
1297 ut.current_date_and_time()
1298 print("*** Starting precon3d.aligner *** \n")
1300 with Progress(
1301 SpinnerColumn(),
1302 TextColumn("[progress.description]{task.description}"),
1303 transient=True,
1304 ) as progress:
1305 task = progress.add_task(description="Aligning...", total=None)
1306 # parse and run prep
1307 if alignment_attrs.registration_method == "turboreg":
1308 fiji_attrs = FijiAttrs(**config_dict["fiji_attrs"])
1309 align_image_stack_turboreg(
1310 general_attrs.input_directory, alignment_attrs, fiji_attrs
1311 )
1312 elif alignment_attrs.registration_method == "sift":
1313 align_image_stack_sift(
1314 general_attrs.input_directory, alignment_attrs
1315 )
1316 elif alignment_attrs.registration_method == "phasecorr":
1317 align_image_stack_phasecorr(
1318 general_attrs.input_directory, alignment_attrs
1319 )
1320 progress.update(task, completed=True) # Mark the task as completed
1322 print("Done!\n")
1324 toc = time.perf_counter()
1326 print(
1327 f"""Data (pre)paration/(con)truction in 3D took
1328 {toc - tic:0.2f} seconds
1329 {(toc - tic)/60:0.2f} minutes
1330 {(toc - tic)/3600:0.2f} hours"""
1331 )
1334@app.callback()
1335def callback():
1336 """
1337 precon3d.aligner aligns image stacks.
1338 """