zstack_registration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib widget\n",
    "\n",
    "from time import time\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import tifffile\n",
    "import concurrent.futures\n",
    "from typing import List, Tuple, Dict\n",
    "import napari\n",
    "import SimpleITK as sitk\n",
    "\n",
    "from vrAnalysis import fileManagement as fm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_organize_stack(folder_path: str, num_planes: int, flyback=0, pattern: str = \"*_00001_*\") -> List[np.ndarray]:\n",
    "    \"\"\"\n",
    "    Load and organize a multi-channel z-stack from tiff files.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    folder_path : str\n",
    "        Path to folder containing tiff files\n",
    "    num_planes: int\n",
    "        Number of planes per stack\n",
    "    flyback: int, optional\n",
    "        Number of flyback planes, if more than 0 will be removed from head of stack\n",
    "    pattern : str, optional\n",
    "        Glob pattern to match relevant files\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    stacks : List[np.ndarray]\n",
    "        List of 3D arrays of data from each color channel (z, y, x)\n",
    "    \"\"\"\n",
    "    folder = Path(folder_path)\n",
    "    files = sorted(folder.glob(pattern))\n",
    "    \n",
    "    if not files:\n",
    "        raise ValueError(f\"No files found matching pattern {pattern}\")\n",
    "    \n",
    "    # Load all tiff files in parallel\n",
    "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "        tiff_data = list(executor.map(tifffile.imread, files))\n",
    "    \n",
    "    # Read first file to get dimensions\n",
    "    if tiff_data[0].ndim == 3:\n",
    "        n_channels = 1\n",
    "        height, width = tiff_data[0].shape[1:]\n",
    "    else:\n",
    "        n_channels, height, width = tiff_data[0].shape[1:]\n",
    "\n",
    "    # Preallocate list for storing data from each channel \n",
    "    channels = [[] for _ in range(n_channels)]\n",
    "\n",
    "    for data in tiff_data:\n",
    "        if n_channels == 1:\n",
    "            data = data.reshape(data.shape[0], 1, height, width)\n",
    "        \n",
    "        for ichannel in range(n_channels):\n",
    "            channels[ichannel].append(data[:, ichannel])\n",
    "\n",
    "        data = None\n",
    "    \n",
    "    # Combine all files\n",
    "    channels = [np.concatenate(channel, axis=0) for channel in channels]\n",
    "    num_extra = channels[0].shape[0] % num_planes\n",
    "    channels = [np.concatenate((channel, np.full((num_planes - num_extra, *channel.shape[1:]), np.nan)), axis=0) for channel in channels]\n",
    "    num_repeats = channels[0].shape[0] // num_planes\n",
    "\n",
    "    # Reshape by plane\n",
    "    stacks = [np.nanmean(channel.reshape(num_repeats, num_planes, *channel.shape[1:]), axis=0) for channel in channels]\n",
    "    \n",
    "    # Remove flyback planes\n",
    "    if flyback > 0:\n",
    "        stacks = [stack[flyback:] for stack in stacks]\n",
    "        \n",
    "    return stacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "from typing import Tuple, Optional, Dict, Any\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "\n",
    "class MetricType(Enum):\n",
    "    MUTUAL_INFO = 'mi'\n",
    "    CORRELATION = 'nc'\n",
    "    MEAN_SQUARES = 'ms'\n",
    "\n",
    "def register_translation(\n",
    "    moving: np.ndarray,\n",
    "    fixed: np.ndarray,\n",
    "    verbose: bool = False\n",
    ") -> Tuple[np.ndarray, sitk.Transform]:\n",
    "    \"\"\"\n",
    "    Register two 3D image stacks using translation only.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Moving image stack (z, y, x)\n",
    "    fixed : np.ndarray\n",
    "        Fixed image stack (z, y, x)\n",
    "    verbose : bool\n",
    "        Print progress\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    registered : np.ndarray\n",
    "        Registered moving image\n",
    "    transform : sitk.Transform\n",
    "        The computed transform\n",
    "    \"\"\"\n",
    "    # Convert to SimpleITK images\n",
    "    moving_sitk = sitk.GetImageFromArray(moving.astype(np.float32))\n",
    "    fixed_sitk = sitk.GetImageFromArray(fixed.astype(np.float32))\n",
    "    \n",
    "    # Initialize registration\n",
    "    registration = sitk.ImageRegistrationMethod()\n",
    "    \n",
    "    # Translation transform\n",
    "    transform = sitk.TranslationTransform(3)\n",
    "    registration.SetInitialTransform(transform)\n",
    "    \n",
    "    # Simple normalized correlation metric\n",
    "    registration.SetMetricAsCorrelation()\n",
    "    \n",
    "    # Simple optimizer\n",
    "    registration.SetOptimizerAsGradientDescent(\n",
    "        learningRate=1.0,\n",
    "        numberOfIterations=100\n",
    "    )\n",
    "    \n",
    "    # Linear interpolation\n",
    "    registration.SetInterpolator(sitk.sitkLinear)\n",
    "    \n",
    "    if verbose:\n",
    "        def callback(filter):\n",
    "            print(f\"Iteration: {filter.GetOptimizerIteration()}, \"\n",
    "                  f\"Metric: {filter.GetMetricValue():.4f}\")\n",
    "        registration.AddCommand(sitk.sitkIterationEvent,\n",
    "                              lambda: callback(registration))\n",
    "    \n",
    "    # Perform registration\n",
    "    final_transform = registration.Execute(fixed_sitk, moving_sitk)\n",
    "    \n",
    "    # Apply transform\n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetReferenceImage(fixed_sitk)\n",
    "    resampler.SetInterpolator(sitk.sitkLinear)\n",
    "    resampler.SetTransform(final_transform)\n",
    "    \n",
    "    registered_sitk = resampler.Execute(moving_sitk)\n",
    "    registered = sitk.GetArrayFromImage(registered_sitk)\n",
    "    \n",
    "    if verbose:\n",
    "        # Print the final translation\n",
    "        translation = final_transform.GetParameters()\n",
    "        print(f\"\\nFinal translation (x,y,z): {translation}\")\n",
    "    \n",
    "    return registered, final_transform\n",
    "\n",
    "\n",
    "\n",
    "def register_constrained_rigid(\n",
    "    moving: np.ndarray,\n",
    "    fixed: np.ndarray,\n",
    "    max_rotation_degrees: float = 15.0,  # Maximum rotation allowed in any direction\n",
    "    translation_scales: Tuple[float, float, float] = (1.0, 1.0, 0.5),  # Scale factors for x,y,z translation\n",
    "    verbose: bool = False\n",
    ") -> Tuple[np.ndarray, sitk.Transform]:\n",
    "    \"\"\"\n",
    "    Register two 3D image stacks using constrained rigid transformation.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Moving image stack (z, y, x)\n",
    "    fixed : np.ndarray\n",
    "        Fixed image stack (z, y, x)\n",
    "    max_rotation_degrees : float\n",
    "        Maximum allowed rotation in any direction (degrees)\n",
    "    translation_scales : tuple\n",
    "        Scale factors for translation optimization (x,y,z)\n",
    "        Use smaller values for z if z-resolution is different\n",
    "    verbose : bool\n",
    "        Print progress\n",
    "    \"\"\"\n",
    "    # Convert to SimpleITK images\n",
    "    moving_sitk = sitk.GetImageFromArray(moving.astype(np.float32))\n",
    "    fixed_sitk = sitk.GetImageFromArray(fixed.astype(np.float32))\n",
    "    \n",
    "    # Initialize registration\n",
    "    registration = sitk.ImageRegistrationMethod()\n",
    "    \n",
    "    # Euler3D transform (rotation + translation)\n",
    "    transform = sitk.Euler3DTransform()\n",
    "    registration.SetInitialTransform(transform)\n",
    "    \n",
    "    # Set transform parameter scaling to constrain rotation\n",
    "    # Order is: rotX, rotY, rotZ, transX, transY, transZ\n",
    "    max_rot_rad = np.deg2rad(max_rotation_degrees)\n",
    "    scales = [1/max_rot_rad]*3 + [1.0/s for s in translation_scales]\n",
    "    registration.SetOptimizerScales(scales)\n",
    "    \n",
    "    # Regular step gradient descent with small steps for rotation\n",
    "    registration.SetOptimizerAsRegularStepGradientDescent(\n",
    "        learningRate=0.1,\n",
    "        minStep=0.0001,\n",
    "        numberOfIterations=200,\n",
    "        relaxationFactor=0.5,\n",
    "        gradientMagnitudeTolerance=1e-6\n",
    "    )\n",
    "    \n",
    "    # Normalized correlation metric\n",
    "    registration.SetMetricAsCorrelation()\n",
    "    \n",
    "    # Linear interpolation\n",
    "    registration.SetInterpolator(sitk.sitkLinear)\n",
    "    \n",
    "    # Multi-resolution approach\n",
    "    registration.SetShrinkFactorsPerLevel([4, 2, 1])\n",
    "    registration.SetSmoothingSigmasPerLevel([2, 1, 0])\n",
    "    \n",
    "    if verbose:\n",
    "        def callback(filter):\n",
    "            if filter.GetOptimizerIteration() == 0:\n",
    "                print(\"\\nStarting registration...\")\n",
    "            print(f\"Iteration: {filter.GetOptimizerIteration()}, \"\n",
    "                  f\"Metric: {filter.GetMetricValue():.4f}\")\n",
    "            # Check if rotations are within bounds\n",
    "            params = filter.GetOptimizerPosition()\n",
    "            rotations_deg = np.rad2deg(params[:3])\n",
    "            if np.any(np.abs(rotations_deg) > max_rotation_degrees):\n",
    "                print(\"Warning: Rotation exceeding specified maximum!\")\n",
    "        registration.AddCommand(sitk.sitkIterationEvent,\n",
    "                              lambda: callback(registration))\n",
    "    \n",
    "    # Perform registration\n",
    "    final_transform = registration.Execute(fixed_sitk, moving_sitk)\n",
    "    \n",
    "    if verbose:\n",
    "        # Print the final parameters\n",
    "        params = final_transform.GetParameters()\n",
    "        rotations_deg = np.rad2deg(params[:3])\n",
    "        translations = params[3:]\n",
    "        print(\"\\nFinal transform parameters:\")\n",
    "        print(f\"Rotations (deg) - X: {rotations_deg[0]:.2f}, \"\n",
    "              f\"Y: {rotations_deg[1]:.2f}, Z: {rotations_deg[2]:.2f}\")\n",
    "        print(f\"Translations - X: {translations[0]:.2f}, \"\n",
    "              f\"Y: {translations[1]:.2f}, Z: {translations[2]:.2f}\")\n",
    "    \n",
    "    # Apply transform\n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetReferenceImage(fixed_sitk)\n",
    "    resampler.SetInterpolator(sitk.sitkLinear)\n",
    "    resampler.SetTransform(final_transform)\n",
    "    \n",
    "    registered_sitk = resampler.Execute(moving_sitk)\n",
    "    registered = sitk.GetArrayFromImage(registered_sitk)\n",
    "    \n",
    "    return registered, final_transform\n",
    "\n",
    "\n",
    "\n",
    "def register_zstacks(\n",
    "    moving: np.ndarray,\n",
    "    fixed: np.ndarray,\n",
    "    transform_type: str = 'rigid',\n",
    "    metric: MetricType = MetricType.CORRELATION,\n",
    "    params: Optional[Dict[str, Any]] = None,\n",
    "    verbose: bool = False\n",
    ") -> Tuple[np.ndarray, sitk.Transform]:\n",
    "    \"\"\"\n",
    "    Register two 3D image stacks using SimpleITK.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Moving image stack (z, y, x)\n",
    "    fixed : np.ndarray\n",
    "        Fixed image stack (z, y, x)\n",
    "    transform_type : str\n",
    "        Type of registration: 'rigid', 'affine', or 'rigid+nonrigid'\n",
    "    metric : MetricType\n",
    "        Similarity metric to use\n",
    "    params : dict, optional\n",
    "        Override default parameters for registration\n",
    "    verbose : bool\n",
    "        Whether to print registration progress\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    registered : np.ndarray\n",
    "        Registered moving image\n",
    "    transform : sitk.Transform\n",
    "        The computed transform(s)\n",
    "    \"\"\"\n",
    "    # Set default parameters\n",
    "    default_params = {\n",
    "        'shrink_factors': (4, 2, 1),\n",
    "        'smooth_sigmas': (2, 1, 0),\n",
    "        'learning_rate': 0.1,\n",
    "        'iterations': 200,\n",
    "        'grad_tol': 1e-6,\n",
    "        'sampling_percentage': 0.01\n",
    "    }\n",
    "    if params:\n",
    "        default_params.update(params)\n",
    "    params = default_params\n",
    "\n",
    "    # Convert to SimpleITK images and normalize\n",
    "    moving_sitk = sitk.GetImageFromArray(moving.astype(np.float32))\n",
    "    fixed_sitk = sitk.GetImageFromArray(fixed.astype(np.float32))\n",
    "    moving_sitk = sitk.Normalize(moving_sitk)\n",
    "    fixed_sitk = sitk.Normalize(fixed_sitk)\n",
    "\n",
    "    if transform_type == 'rigid+nonrigid':\n",
    "        # First do rigid registration\n",
    "        registered_sitk, rigid_transform = _register_single_transform(\n",
    "            moving_sitk, fixed_sitk,\n",
    "            transform='rigid',\n",
    "            metric=metric,\n",
    "            params=params,\n",
    "            verbose=verbose\n",
    "        )\n",
    "        # Then do nonrigid registration starting from rigid result\n",
    "        registered_sitk, bspline_transform = _register_single_transform(\n",
    "            registered_sitk, fixed_sitk,\n",
    "            transform='bspline',\n",
    "            metric=metric,\n",
    "            params=params,\n",
    "            verbose=verbose\n",
    "        )\n",
    "        # Combine transforms\n",
    "        composite_transform = sitk.CompositeTransform([bspline_transform, rigid_transform])\n",
    "        return sitk.GetArrayFromImage(registered_sitk), composite_transform\n",
    "    else:\n",
    "        # Single transform registration\n",
    "        registered_sitk, transform = _register_single_transform(\n",
    "            moving_sitk, fixed_sitk,\n",
    "            transform=transform_type,\n",
    "            metric=metric,\n",
    "            params=params,\n",
    "            verbose=verbose\n",
    "        )\n",
    "        return sitk.GetArrayFromImage(registered_sitk), transform\n",
    "\n",
    "def _register_single_transform(\n",
    "    moving_sitk: sitk.Image,\n",
    "    fixed_sitk: sitk.Image,\n",
    "    transform: str,\n",
    "    metric: MetricType,\n",
    "    params: Dict[str, Any],\n",
    "    verbose: bool\n",
    ") -> Tuple[sitk.Image, sitk.Transform]:\n",
    "    \"\"\"Helper function for single-transform registration.\"\"\"\n",
    "    \n",
    "    registration = sitk.ImageRegistrationMethod()\n",
    "    \n",
    "    # Set up transform\n",
    "    if transform == 'rigid':\n",
    "        initial_transform = sitk.Euler3DTransform()\n",
    "    elif transform == 'affine':\n",
    "        initial_transform = sitk.AffineTransform(3)\n",
    "    elif transform == 'bspline':\n",
    "        transform_domain_mesh_size = [8] * 3\n",
    "        initial_transform = sitk.BSplineTransformInitializer(\n",
    "            fixed_sitk, transform_domain_mesh_size\n",
    "        )\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown transform type: {transform}\")\n",
    "    \n",
    "    registration.SetInitialTransform(initial_transform)\n",
    "    \n",
    "    # Set up metric\n",
    "    if metric == MetricType.MUTUAL_INFO:\n",
    "        registration.SetMetricAsMattesMutualInformation()\n",
    "    elif metric == MetricType.CORRELATION:\n",
    "        registration.SetMetricAsCorrelation()\n",
    "    elif metric == MetricType.MEAN_SQUARES:\n",
    "        registration.SetMetricAsMeanSquares()\n",
    "    \n",
    "    # Set up optimizer\n",
    "    if transform == 'bspline':\n",
    "        registration.SetOptimizerAsLBFGSB(\n",
    "            gradientConvergenceTolerance=params['grad_tol'],\n",
    "            numberOfIterations=params['iterations'],\n",
    "            maximumNumberOfCorrections=5,\n",
    "            maximumNumberOfFunctionEvaluations=1000,\n",
    "            costFunctionConvergenceFactor=1e7\n",
    "        )\n",
    "    else:\n",
    "        registration.SetOptimizerAsGradientDescent(\n",
    "            learningRate=params['learning_rate'],\n",
    "            numberOfIterations=params['iterations'],\n",
    "            convergenceMinimumValue=params['grad_tol'],\n",
    "            convergenceWindowSize=10\n",
    "        )\n",
    "    \n",
    "    # Set up sampling\n",
    "    registration.SetMetricSamplingPercentage(params['sampling_percentage'])\n",
    "    registration.SetMetricSamplingStrategy(registration.RANDOM)\n",
    "    \n",
    "    # Set up interpolator\n",
    "    registration.SetInterpolator(sitk.sitkLinear)\n",
    "    \n",
    "    # Set up multi-resolution framework\n",
    "    registration.SetShrinkFactorsPerLevel(params['shrink_factors'])\n",
    "    registration.SetSmoothingSigmasPerLevel(params['smooth_sigmas'])\n",
    "    \n",
    "    if verbose:\n",
    "        def callback(filter):\n",
    "            print(f\"Iteration: {filter.GetOptimizerIteration()}, \"\n",
    "                  f\"Metric: {filter.GetMetricValue():.4f}\")\n",
    "        registration.AddCommand(sitk.sitkIterationEvent,\n",
    "                              lambda: callback(registration))\n",
    "    \n",
    "    # Perform registration\n",
    "    final_transform = registration.Execute(fixed_sitk, moving_sitk)\n",
    "    \n",
    "    # Apply transform\n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetReferenceImage(fixed_sitk)\n",
    "    resampler.SetInterpolator(sitk.sitkLinear)\n",
    "    resampler.SetTransform(final_transform)\n",
    "    \n",
    "    return resampler.Execute(moving_sitk), final_transform\n",
    "\n",
    "\n",
    "def visualize_registration(moving: np.ndarray, fixed: np.ndarray, registered: np.ndarray) -> None:\n",
    "    \"\"\"\n",
    "    Visualize registration results using napari.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Original moving image\n",
    "    fixed : np.ndarray\n",
    "        Fixed image\n",
    "    registered : np.ndarray\n",
    "        Registered moving image\n",
    "    \"\"\"\n",
    "    viewer = napari.Viewer()\n",
    "    \n",
    "    # Add layers\n",
    "    viewer.add_image(fixed, name='Fixed', colormap='green', blending='additive')\n",
    "    viewer.add_image(registered, name='Registered', colormap='red', blending='additive')\n",
    "    viewer.add_image(moving, name='Original Moving', colormap='orange', blending='additive', visible=False)\n",
    "    \n",
    "    return viewer\n",
    "\n",
    "def compare_metrics(moving: np.ndarray, fixed: np.ndarray, metrics: list = ['mi', 'nc', 'ms'], verbose=False) -> Dict[str, Tuple[np.ndarray, sitk.Transform]]:\n",
    "    \"\"\"\n",
    "    Compare different registration metrics.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Moving image stack\n",
    "    fixed : np.ndarray\n",
    "        Fixed image stack\n",
    "    metrics : list\n",
    "        List of metrics to try\n",
    "    verbose : bool\n",
    "        Whether to print registration progress\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    dict\n",
    "        Dictionary of results for each metric\n",
    "    \"\"\"\n",
    "    results = {}\n",
    "    \n",
    "    for metric in metrics:\n",
    "        print(f\"Trying {metric} metric...\")\n",
    "        registered, transform = register_zstacks(moving, fixed, metric=metric, verbose=verbose)\n",
    "        results[metric] = (registered, transform)\n",
    "        \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import Tuple\n",
    "from scipy import fftpack\n",
    "\n",
    "def phase_correlation_3d(\n",
    "    moving: np.ndarray,\n",
    "    fixed: np.ndarray,\n",
    "    upsample_factor: int = 1\n",
    ") -> Tuple[np.ndarray, Tuple[int, int, int]]:\n",
    "    \"\"\"\n",
    "    Register two 3D image stacks using phase correlation.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    moving : np.ndarray\n",
    "        Moving image stack (z, y, x)\n",
    "    fixed : np.ndarray\n",
    "        Fixed image stack (z, y, x)\n",
    "    upsample_factor : int\n",
    "        Factor for subpixel accuracy. 1 means pixel accuracy.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    registered : np.ndarray\n",
    "        Registered moving image\n",
    "    shifts : tuple\n",
    "        (z, y, x) shifts needed to register images\n",
    "    \"\"\"\n",
    "    # Ensure float arrays\n",
    "    moving = moving.astype(float)\n",
    "    fixed = fixed.astype(float)\n",
    "    \n",
    "    # Compute FFTs\n",
    "    fixed_fft = fftpack.fftn(fixed)\n",
    "    moving_fft = fftpack.fftn(moving)\n",
    "    \n",
    "    # Compute cross-power spectrum\n",
    "    cross_power = moving_fft * fixed_fft.conjugate()\n",
    "    \n",
    "    # Normalize to get only phase information\n",
    "    eps = np.finfo(float).eps\n",
    "    normalized_cross_power = cross_power / (np.abs(cross_power) + eps)\n",
    "    \n",
    "    # Compute phase correlation\n",
    "    correlation = np.real(fftpack.ifftn(normalized_cross_power))\n",
    "    \n",
    "    # Find shift that gives maximum correlation\n",
    "    shifts = np.unravel_index(np.argmax(correlation), correlation.shape)\n",
    "    \n",
    "    # Convert shifts to centered coordinates\n",
    "    shifts = [\n",
    "        s if s <= shape//2 else s - shape \n",
    "        for s, shape in zip(shifts, correlation.shape)\n",
    "    ]\n",
    "    \n",
    "    # Apply shift to register images\n",
    "    registered = np.roll(np.roll(np.roll(\n",
    "        moving,\n",
    "        shifts[0], axis=0),\n",
    "        shifts[1], axis=1),\n",
    "        shifts[2], axis=2)\n",
    "    \n",
    "    return registered, tuple(shifts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "num_planes = 30\n",
    "flyback = 4\n",
    "folder_path_0 = fm.serverPath(zaru=False) / \"ATL060\" / \"2024-07-12\" / \"702\"\n",
    "folder_path_1 = fm.serverPath(zaru=False) / \"ATL060\" / \"2024-07-15\" / \"702\"\n",
    "\n",
    "# get stacks\n",
    "stacks_0 = load_and_organize_stack(folder_path_0, num_planes, flyback=flyback)\n",
    "stacks_1 = load_and_organize_stack(folder_path_1, num_planes, flyback=flyback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Translations:\n",
      "X: 0.023\n",
      "Y: 0.017\n",
      "Z: -0.001\n",
      "Detected shifts (z,y,x): (2, 1, 0)\n"
     ]
    }
   ],
   "source": [
    "# Load and preprocess stacks\n",
    "moving_stack = stacks_0[1]\n",
    "fixed_stack = stacks_1[1]   \n",
    "\n",
    "# Try registration\n",
    "registered, transform = register_constrained_rigid(moving_stack, fixed_stack, max_rotation_degrees=5.0, translation_scales=(1.0, 1.0, 0.5), verbose=False)\n",
    "params = transform.GetParameters()\n",
    "translations = params[3:]\n",
    "print(\"Translations:\")\n",
    "print(f\"X: {translations[0]:.3f}\")\n",
    "print(f\"Y: {translations[1]:.3f}\")\n",
    "print(f\"Z: {translations[2]:.3f}\")\n",
    "\n",
    "registered, shifts = phase_correlation_3d(moving_stack, fixed_stack)\n",
    "print(f\"Detected shifts (z,y,x): {shifts}\")\n",
    "\n",
    "viewer = visualize_registration(moving_stack, fixed_stack, registered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For rigid registration\n",
    "registered, transform = register_zstacks(stacks_0[1], stacks_1[1], transform_type='rigid', verbose=False)\n",
    "\n",
    "# # For rigid + nonrigid registration\n",
    "# registered, transform = register_zstacks(\n",
    "#     moving_stack, \n",
    "#     fixed_stack,\n",
    "#     transform_type='rigid+nonrigid',\n",
    "#     metric=MetricType.CORRELATION,\n",
    "#     params={'learning_rate': 0.05, 'iterations': 300},\n",
    "#     verbose=True\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Trying mi metric...\n",
      "\n",
      "Trying nc metric...\n",
      "\n",
      "Trying ms metric...\n"
     ]
    }
   ],
   "source": [
    "results = compare_metrics(stacks_0[1], stacks_1[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = visualize_registration(stacks_0[1], stacks_1[1], stacks_0[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# register stacks to each other\n",
    "best_metric = 'nc'  # or whichever you found works best\n",
    "best_channel = 0 # green to start\n",
    "registered, transform = register_zstacks(stacks_0[best_channel], stacks_1[best_channel], metric=best_metric, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = visualize_registration(stacks_0[best_channel], stacks_1[best_channel], registered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import napari\n",
    "\n",
    "# view stacks[0] in napari\n",
    "viewer = napari.Viewer()\n",
    "viewer.add_image(stacks_1[0], name='Channel 0', colormap='green', blending='additive')\n",
    "viewer.add_image(stacks_1[1], name='Channel 1', colormap='red', blending='additive')\n",
    "\n",
    "# show viewer\n",
    "viewer.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from imageio.v3 import imread\n",
    "import concurrent.futures\n",
    "\n",
    "def benchmark_loading(folder_path: str, pattern: str = \"*_00001_*\"):\n",
    "    files = sorted(Path(folder_path).glob(pattern))\n",
    "    \n",
    "    # Test tifffile\n",
    "    start = time.time()\n",
    "    a = [tifffile.imread(f) for f in files]\n",
    "    print(f\"tifffile sequential: {time.time() - start:.2f}s\")\n",
    "    \n",
    "    # Test tifffile with threading\n",
    "    start = time.time()\n",
    "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "        b = list(executor.map(tifffile.imread, files))\n",
    "    print(f\"tifffile threaded: {time.time() - start:.2f}s\")\n",
    "    \n",
    "    # Test imageio\n",
    "    start = time.time()\n",
    "    c = [imread(f) for f in files]\n",
    "    print(f\"imageio sequential: {time.time() - start:.2f}s\")\n",
    "\n",
    "    # Test tifffile again\n",
    "    start = time.time()\n",
    "    d = [tifffile.imread(f) for f in files]\n",
    "    print(f\"tifffile sequential again: {time.time() - start:.2f}s\")\n",
    "\n",
    "    print(all([np.allclose(x, y) for x, y in zip(a, b)]))\n",
    "    print(all([np.allclose(x, y) for x, y in zip(a, c)]))\n",
    "    print(all([np.allclose(x, y) for x, y in zip(a, d)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tifffile sequential: 16.22s\n",
      "tifffile threaded: 0.97s\n",
      "imageio sequential: 1.52s\n",
      "tifffile sequential again: 1.36s\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "benchmark_loading(folder_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vrAnalysis",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}