import torch as t
import torch.nn.functional as tfun
import torchvision.transforms as trans
import torch.optim as topt
import torch as tc
import torch.nn.functional as F
from typing import Optional, Tuple, List, Union, Dict, Any
import numpy as np
import math
import cv2
from skimage import color
from pathlib import Path
from tqdm.auto import tqdm
[docs]
def build_reference_coordinate_system(input_tensor: Optional[t.Tensor] = None,
dimensions: Optional[t.Size] = None,
compute_device: Optional[Union[str, t.device]] = None) -> t.Tensor:
if input_tensor is not None:
dimensions = input_tensor.size()
# Convert string device specification to torch.device
if isinstance(compute_device, str):
compute_device = t.device(compute_device)
if compute_device is None and input_tensor is not None:
base_transform = t.eye(len(dimensions)-1)[:-1, :].unsqueeze(0).type_as(input_tensor)
else:
base_transform = t.eye(len(dimensions)-1, device=compute_device)[:-1, :].unsqueeze(0)
base_transform = t.repeat_interleave(base_transform, dimensions[0], dim=0)
coordinate_grid = tfun.affine_grid(base_transform, dimensions, align_corners=False)
return coordinate_grid
[docs]
def gaussian_smoothing(input_tensor: t.Tensor, blur_sigma: float) -> t.Tensor:
with t.set_grad_enabled(False):
kernel_width = int(blur_sigma * 2.54) + 1
if kernel_width % 2 == 0:
kernel_width += 1
return trans.GaussianBlur(kernel_width, blur_sigma)(input_tensor)
[docs]
def scale_tensor_to_dimensions(input_tensor: t.Tensor,
target_dimensions: t.Size,
interpolation_method: str = 'bilinear') -> t.Tensor:
return tfun.interpolate(input_tensor, size=target_dimensions,
mode=interpolation_method, align_corners=False)
[docs]
def compute_normalized_cross_correlation(sources: t.Tensor,
targets: t.Tensor,
device: Optional[Union[str, t.device]] = None,
**config_params) -> t.Tensor:
ndim = len(sources.size()) - 2
if ndim not in [2, 3]:
raise ValueError("Unsupported number of dimensions.")
try:
size =7
except:
size = 3
window = (size, ) * ndim
if device is None:
sum_filt = tc.ones([1, 1, *window]).type_as(sources)
else:
sum_filt = tc.ones([1, 1, *window], device=device)
pad_no = math.floor(window[0] / 2)
stride = ndim * (1,)
padding = ndim * (pad_no,)
conv_fn = getattr(F, 'conv%dd' % ndim)
sources_denom = sources**2
targets_denom = targets**2
numerator = sources*targets
sources_sum = conv_fn(sources, sum_filt, stride=stride, padding=padding)
targets_sum = conv_fn(targets, sum_filt, stride=stride, padding=padding)
sources_denom_sum = conv_fn(sources_denom, sum_filt, stride=stride, padding=padding)
targets_denom_sum = conv_fn(targets_denom, sum_filt, stride=stride, padding=padding)
numerator_sum = conv_fn(numerator, sum_filt, stride=stride, padding=padding)
size = np.prod(window)
u_sources = sources_sum / size
u_targets = targets_sum / size
cross = numerator_sum - u_targets * sources_sum - u_sources * targets_sum + u_sources * u_targets * size
sources_var = sources_denom_sum - 2 * u_sources * sources_sum + u_sources * u_sources * size
targets_var = targets_denom_sum - 2 * u_targets * targets_sum + u_targets * u_targets * size
ncc = cross * cross / (sources_var * targets_var + 1e-5)
return -tc.mean(ncc)
[docs]
def create_multiscale_representation(input_tensor: t.Tensor,
level_count: int,
interpolation_method: str = 'bilinear',
scale_factor: float = 2.0) -> List[t.Tensor]:
pyramid_levels = [None] * level_count
# Build from fine to coarse
for i in range(level_count - 1, -1, -1):
if i == level_count - 1:
# Original resolution
pyramid_levels[i] = input_tensor
else:
# Get previous level and compute dimensions for current level
prev_size = pyramid_levels[i+1].size()
current_dims = tuple(int(prev_size[j] / scale_factor) if j > 1 else prev_size[j]
for j in range(len(prev_size)))
# Extract just the spatial dimensions
spatial_dims = t.Size(current_dims)[2:]
# Apply smoothing to prevent aliasing, then downsample
smoothed = gaussian_smoothing(pyramid_levels[i+1], 1)
downsampled = scale_tensor_to_dimensions(smoothed, spatial_dims,
interpolation_method)
pyramid_levels[i] = downsampled
return pyramid_levels
[docs]
def convert_image_to_tensor(img_array: np.ndarray, compute_device: Union[str, t.device] = "cpu") -> t.Tensor:
# Convert string device specification to torch.device
if isinstance(compute_device, str):
compute_device = t.device(compute_device)
# Normalize image if it's not already in [0, 1] range
if img_array.dtype != np.float32 and img_array.dtype != np.float64:
if img_array.max() > 1.0:
img_array = img_array.astype(np.float32) / 255.0
if len(img_array.shape) == 3:
# Color image
return t.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0).to(compute_device)
elif len(img_array.shape) == 2:
# Grayscale image
return t.from_numpy(img_array).unsqueeze(0).unsqueeze(0).to(compute_device)
else:
raise ValueError(f"Unsupported image dimensions: {img_array.shape}")
[docs]
def prepare_image_tensors(source_image: np.ndarray,
target_image: np.ndarray,
compute_device: Union[str, t.device],
normalize: bool = True) -> Tuple[t.Tensor, t.Tensor]:
# Convert string device specification to torch.device
if isinstance(compute_device, str):
compute_device = t.device(compute_device)
# Convert to grayscale if RGB
if len(source_image.shape) == 3 and source_image.shape[2] == 3:
gray_source = color.rgb2gray(source_image)
else:
gray_source = source_image
if len(target_image.shape) == 3 and target_image.shape[2] == 3:
gray_target = color.rgb2gray(target_image)
else:
gray_target = target_image
# Normalize if requested
if normalize:
gray_source = (gray_source - gray_source.min()) / (gray_source.max() - gray_source.min() + 1e-10)
gray_target = (gray_target - gray_target.min()) / (gray_target.max() - gray_target.min() + 1e-10)
# Convert to tensor format
tensor_source = convert_image_to_tensor(gray_source, compute_device)
tensor_target = convert_image_to_tensor(gray_target, compute_device)
# Create tensors with gradient tracking
source_tensor = t.tensor(tensor_source, dtype=t.float32, requires_grad=True).to(compute_device)
target_tensor = t.tensor(tensor_target, dtype=t.float32, requires_grad=True).to(compute_device)
return source_tensor, target_tensor
[docs]
def elastic_image_registration(
source: np.ndarray,
target: np.ndarray,
similarity_metric: str = "ncc",
similarity_metric_params: Dict[str, Any] = {"size": 7},
compute_device: Union[str, t.device] = "cuda",
verbose: bool = False,
output_dir: Optional[Union[str, Path]] = None,
save_intermediate: bool = False,
) -> Tuple[t.Tensor, t.Tensor]:
# Setup
device = t.device(compute_device) if isinstance(compute_device, str) else compute_device
src_t, tgt_t = prepare_image_tensors(source, target, device)
aligned_source = cv2.warpAffine(source, np.eye(2, 3), (target.shape[1], target.shape[0]), borderMode=cv2.BORDER_REFLECT)
source_t, target_t = prepare_image_tensors(aligned_source, target, device)
pyramid_levels = 6
src_pyr = create_multiscale_representation(source_t, pyramid_levels)
tgt_pyr = create_multiscale_representation(target_t, pyramid_levels)
# Hyperparameters
iterations_per_level = [400, 400, 400, 400, 400, 400, 400]
learning_rates = [0.005, 0.0025, 0.0025, 0.0025, 0.0025, 0.0025, 0.0015]
regularization_weights = [1.5, 1.5, 1.5, 1.5, 1.5, 1.2, 0.6]
prev_def_field = None
# Loop through pyramid levels
for lvl in range(pyramid_levels):
curr_src = src_pyr[lvl]
curr_tgt = tgt_pyr[lvl]
H, W = curr_src.shape[2:]
# Initialize or upsample deformation field
if lvl == 0:
def_field = initialize_deformation_field(curr_src).detach().clone().requires_grad_(True)
else:
def_field = scale_deformation_field(prev_def_field, (H, W)).detach().clone().requires_grad_(True)
# Optimizer: LBFGS on final level, Adam otherwise
# if lvl == pyramid_levels - 1:
# optimizer = topt.LBFGS([def_field], lr=learning_rates[lvl], max_iter=50, line_search_fn="strong_wolfe")
# else:
optimizer = topt.Adam([def_field], lr=learning_rates[lvl])
weight = regularization_weights[lvl]
for iter_idx in tqdm(range(iterations_per_level[lvl]), disable=not verbose, desc=f"Level {lvl}/{pyramid_levels-1}"):
def closure():
optimizer.zero_grad()
warped = apply_deformation_field(curr_src, def_field, compute_device=device)
sim_loss = compute_normalized_cross_correlation(warped, curr_tgt, compute_device=device, **similarity_metric_params)
reg_loss = deformation_loss(def_field, compute_device=device)
loss = sim_loss + weight * reg_loss
loss.backward()
return loss
loss = optimizer.step(closure)
with t.no_grad():
# Optional clipping to prevent folding
max_disp = 5.0 # pixels
def_field.clamp_(-max_disp, max_disp)
prev_def_field = def_field
# Upsample to original shape if needed
final_def = scale_deformation_field(prev_def_field, (src_t.size(2), src_t.size(3))) if pyramid_levels != pyramid_levels else prev_def_field
final_warped = apply_deformation_field(src_t, final_def, compute_device=device)
# # Save outputs if needed
# if output_dir:
# os.makedirs(output_dir, exist_ok=True)
# cv2.imwrite(os.path.join(output_dir, "final_warped.png"), (final_warped.detach().cpu().numpy()[0, 0] * 255).astype(np.uint8))
return final_def, final_warped