gevi/reproduction_effort/make_test_data_aligned.py

import torch
import torchvision as tv  # type: ignore
import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio  # type: ignore
import json

from functions.align_cameras import align_cameras

if torch.cuda.is_available():
    device_name: str = "cuda:0"
else:
    device_name = "cpu"
print(f"Using device: {device_name}")
device: torch.device = torch.device(device_name)
dtype: torch.dtype = torch.float32

filename_bin_mat: str = "bin_old/Exp001_Trial001_Part001.mat"
filename_bin_mat_fake: str = "Exp001_Trial001_Part001_fake.mat"
fill_value: float = 0.0

mat_data = torch.tensor(
    sio.loadmat(filename_bin_mat)["nparray"].astype(dtype=np.float32),
    dtype=dtype,
    device=device,
)

angle_refref_target = torch.tensor(
    [2],
    dtype=dtype,
    device=device,
)

tvec_refref_target = torch.tensor(
    [10, 3],
    dtype=dtype,
    device=device,
)


t = (
    torch.arange(
        0,
        mat_data.shape[-2],
        dtype=dtype,
        device=device,
    )
    - mat_data.shape[-2] // 2
) / float(mat_data.shape[-2] // 2)

f_a: float = 8
A_a: float = 2
a_target = A_a * torch.sin(2 * torch.pi * t * f_a)

f_x: float = 5
A_x: float = 10
x_target = A_x * torch.sin(2.0 * torch.pi * t * f_x)

f_y: float = 7
A_y: float = 7
y_target = A_y * torch.sin(2 * torch.pi * t * f_y)

master_images: torch.Tensor = mat_data[:, :, mat_data.shape[-2] // 2, 1]

master_images_2: torch.Tensor = master_images.unsqueeze(-1).tile(
    (1, 1, mat_data.shape[-1])
)

# Rotate and move the whole timeseries of the acceptor and oxygenation
master_images_2[..., 0] = tv.transforms.functional.affine(
    img=master_images_2[..., 0].unsqueeze(0),
    angle=-float(angle_refref_target),
    translate=[0, 0],
    scale=1.0,
    shear=0,
    interpolation=tv.transforms.InterpolationMode.BILINEAR,
    fill=fill_value,
).squeeze(0)

master_images_2[..., 0] = tv.transforms.functional.affine(
    img=master_images_2[..., 0].unsqueeze(0),
    angle=0,
    translate=[tvec_refref_target[1], tvec_refref_target[0]],
    scale=1.0,
    shear=0,
    interpolation=tv.transforms.InterpolationMode.BILINEAR,
    fill=fill_value,
).squeeze(0)

master_images_2[..., 2] = tv.transforms.functional.affine(
    img=master_images_2[..., 2].unsqueeze(0),
    angle=-float(angle_refref_target),
    translate=[0, 0],
    scale=1.0,
    shear=0,
    interpolation=tv.transforms.InterpolationMode.BILINEAR,
    fill=fill_value,
).squeeze(0)

master_images_2[..., 2] = tv.transforms.functional.affine(
    img=master_images_2[..., 2].unsqueeze(0),
    angle=0,
    translate=[tvec_refref_target[1], tvec_refref_target[0]],
    scale=1.0,
    shear=0,
    interpolation=tv.transforms.InterpolationMode.BILINEAR,
    fill=fill_value,
).squeeze(0)

fake_data = master_images_2.unsqueeze(-2).tile((1, 1, mat_data.shape[-2], 1)).clone()

for t_id in range(0, fake_data.shape[-2]):
    for c_id in range(0, fake_data.shape[-1]):
        fake_data[..., t_id, c_id] = tv.transforms.functional.affine(
            img=fake_data[..., t_id, c_id].unsqueeze(0),
            angle=-float(a_target[t_id]),
            translate=[0, 0],
            scale=1.0,
            shear=0,
            interpolation=tv.transforms.InterpolationMode.BILINEAR,
            fill=fill_value,
        ).squeeze(0)

fake_data = fake_data.clone()

for t_id in range(0, fake_data.shape[-2]):
    for c_id in range(0, fake_data.shape[-1]):
        fake_data[..., t_id, c_id] = tv.transforms.functional.affine(
            img=fake_data[..., t_id, c_id].unsqueeze(0),
            angle=0.0,
            translate=[y_target[t_id], x_target[t_id]],
            scale=1.0,
            shear=0,
            interpolation=tv.transforms.InterpolationMode.BILINEAR,
            fill=fill_value,
        ).squeeze(0)

fake_data_np = fake_data.cpu().numpy()
mdic = {"nparray": fake_data_np}
sio.savemat(filename_bin_mat_fake, mdic)

# ----------------------------------------------------

batch_size: int = 200
filename_raw_json: str = "raw/Exp001_Trial001_Part001_meta.txt"

with open(filename_raw_json, "r") as file_handle:
    metadata: dict = json.load(file_handle)
channels: list[str] = metadata["channelKey"]


data = torch.tensor(
    sio.loadmat(filename_bin_mat_fake)["nparray"].astype(np.float32),
    dtype=dtype,
    device=device,
)

ref_image = data[:, :, data.shape[-2] // 2, :].clone()

(
    acceptor,
    donor,
    oxygenation,
    volume,
    angle_donor_volume,
    tvec_donor_volume,
    angle_refref,
    tvec_refref,
) = align_cameras(
    channels=channels,
    data=data,
    ref_image=ref_image,
    device=device,
    dtype=dtype,
    batch_size=batch_size,
    fill_value=-1,
)
del data


print("References Acceptor <-> Donor:")
print("Rotation:")
print(f"target: {float(angle_refref_target):.3f}")
print(f"found: {-float(angle_refref):.3f}")
print("Translation")
print(f"target: {float(tvec_refref_target[0]):.3f}, {float(tvec_refref_target[1]):.3f}")
print(f"found: {-float(tvec_refref[0]):.3f}, {-float(tvec_refref[1]):.3f}")

plt.subplot(3, 1, 1)
plt.plot(-angle_donor_volume.cpu(), "g", label="angle found")
plt.plot(a_target.cpu(), "--k", label="angle target")
plt.legend()

plt.subplot(3, 1, 2)
plt.plot(-tvec_donor_volume[:, 0].cpu(), "g", label="x found")
plt.plot(x_target.cpu(), "k--", label="x target")
plt.legend()

plt.subplot(3, 1, 3)
plt.plot(-tvec_donor_volume[:, 1].cpu(), "g", label="y found")
plt.plot(y_target.cpu(), "k--", label="y target")
plt.legend()

plt.show()