File size: 1,281 Bytes
36c95ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import os
import cv2
import imageio
import torch
import kornia as K
import kornia.geometry as KG
def load_timg(file_name):
"""Loads the image with OpenCV and converts to torch.Tensor."""
assert os.path.isfile(file_name), f"Invalid file {file_name}" # nosec
# load image with OpenCV
img = cv2.imread(file_name, cv2.IMREAD_COLOR)
# convert image to torch tensor
tensor = K.image_to_tensor(img, None).float() / 255.
return K.color.bgr_to_rgb(tensor)
registrator = KG.ImageRegistrator('similarity')
img1 = K.resize(load_timg('/Users/oldufo/datasets/stewart/MR-CT/CT.png'), (400, 600))
img2 = K.resize(load_timg('/Users/oldufo/datasets/stewart/MR-CT/MR.png'), (400, 600))
model, intermediate = registrator.register(img1, img2, output_intermediate_models=True)
video_writer = imageio.get_writer('medical_registration.gif', fps=2)
timg_dst_first = img1.clone()
timg_dst_first[0, 0, :, :] = img2[0, 0, :, :]
video_writer.append_data(K.tensor_to_image((timg_dst_first * 255.).byte()))
with torch.no_grad():
for m in intermediate:
timg_dst = KG.homography_warp(img1, m, img2.shape[-2:])
timg_dst[0, 0, :, :] = img2[0, 0, :, :]
video_writer.append_data(K.tensor_to_image((timg_dst_first * 255.).byte()))
video_writer.close()
|