align_OCT.py

"""
Convert exported, rect OCT images to aligned rect and radial, save as NRRD files.

1. Change the `input_folder` to point to the directory containing images.
2. Change the row range used for correlation by changing `rowStart` and `rowEnd` to make sure they only include the tubing and avoid any tissue in this window. 

Install dependencies:

```
python3 -m pip install numpy matplotlib pillow pynrrd tqdm opencv-python
```
"""
# %%
from pathlib import Path

from PIL import Image
from tqdm import tqdm
import nrrd
import numpy as np
import matplotlib.pyplot as plt

# Change this input folder
input_folder = Path("20250116 endometrial 9/4")

pattern = "**/rect-*.tiff"
output_file = input_folder / "rect-align.nrrd"

output_folder = input_folder.parent / (input_folder.stem + "_align")


# %%
def load_images(folder: Path) -> np.ndarray:
    """Load images from a folder and return a 3D np.ndarray"""
    dtype = np.uint8

    image_files = sorted(folder.glob(pattern))
    if not image_files:
        raise ValueError(f"No files found in folder {folder}")

    first_image = Image.open(image_files[0])
    width, height = first_image.size
    width //= 2
    height //= 2

    image_stack = np.zeros((len(image_files), height, width), dtype=dtype)

    for i, image_file in enumerate(tqdm(image_files)):
        image = Image.open(image_file)
        image = image.resize((width, height))
        image_stack[i] = np.array(image)

    return image_stack


# %%
image_stack = load_images(input_folder)

# %%
### Decide which rows to use for alignment.
# Pick just the tubing
image1 = image_stack[101]
rowStart = 70
rowEnd = 110
plt.imshow(image1[rowStart:rowEnd], "gray")

# %%
import cv2

image_stack_align = np.zeros_like(image_stack)

lastImage = image_stack[0].astype(np.float32)
allCorr = np.zeros(len(image_stack))
for i in tqdm(range(0, len(image_stack))):
    image = image_stack[i].astype(np.float32)

    (x, y), corr = cv2.phaseCorrelate(
        image[rowStart:rowEnd], lastImage[rowStart:rowEnd]
    )
    # print(x, y, corr)

    # Use the last aligned image if correlation drops below 0.2
    if corr < 0.2:
        lastImage = imageAlign
        (x, y), corr = cv2.phaseCorrelate(
            image[rowStart:rowEnd], lastImage[rowStart:rowEnd]
        )

    imageAlign = np.roll(image, round(x))
    image_stack_align[i] = imageAlign
    allCorr[i] = corr

# %%
plt.plot(allCorr)
plt.plot(np.zeros_like(allCorr))

# %%
# %matplotlib qt
plt.subplot(311)
plt.imshow(lastImage[rowStart:rowEnd])
plt.subplot(312)
plt.imshow(image[rowStart:rowEnd])
plt.subplot(313)
imageAlign = np.roll(image, round(x))
plt.imshow(imageAlign[rowStart:rowEnd])

plt.tight_layout()


# %%
def show_pair(image1, image2):
    strip1 = image[rowStart:rowEnd].astype(np.float32)
    strip2 = lastImage[rowStart:rowEnd].astype(np.float32)
    (x, y), corr = cv2.phaseCorrelate(strip1, strip2)
    plt.subplot(411)
    plt.imshow(image1, "gray")
    plt.title(f"({x:.2f}, {y:.2f}), {corr:.3f}")
    plt.subplot(412)
    plt.imshow(image2, "gray")
    plt.subplot(413)
    plt.imshow(strip1, "gray")
    plt.subplot(414)
    plt.imshow(strip2, "gray")

    plt.tight_layout()


show_pair(image_stack_align[100], image_stack_align[200])

# %%
print(output_file)
nrrd.write(str(output_file), image_stack_align)


# %%
def makeRadialImage(rectImage, padTop=0):
    dim = min(rectImage.shape[:2])
    dsize = (dim * 2, dim * 2)
    radius = dim
    center = (radius, radius)

    padded = cv2.copyMakeBorder(rectImage, padTop, 0, 0, 0, 0, cv2.BORDER_CONSTANT)
    flags = cv2.WARP_FILL_OUTLIERS + cv2.WARP_INVERSE_MAP
    radialImage = cv2.warpPolar(padded.T, dsize, center, radius, flags)

    radialImage = cv2.flip(radialImage, 1)
    return radialImage


imageRadial = makeRadialImage(image_stack_align[0], 150)
plt.imshow(imageRadial, "gray")
# %%
imageRadial = makeRadialImage(image_stack_align[0], 150)

radial_stack_align = []
for image in tqdm(image_stack_align):
    radial_image = makeRadialImage(image, 150)
    radial_stack_align.append(radial_image)

# %%
radial_stack_align = np.asarray(radial_stack_align)

# %%
output_radial = input_folder / "radial-align.nrrd"
nrrd.write(str(output_radial), radial_stack_align)
# %%