Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 4,035 Bytes
1a2ad19 fbffdcd 1a2ad19 fbffdcd 1a2ad19 ffb825b 1a2ad19 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import numpy as np
import gradio as gr
import imageio
import cv2
import kornia as K
import kornia.geometry as KG
from copy import deepcopy
from tqdm import tqdm
from base64 import b64encode
import torch
import torch.nn.functional as F
use_cuda: bool = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
registrator = KG.ImageRegistrator('similarity',
loss_fn = F.mse_loss,
lr=8e-4, pyramid_levels=3, num_iterations=500).to(device)
models = []
def resize_images(f_names):
for i, f_name in enumerate(f_names):
img = cv2.imread(f_name, cv2.IMREAD_COLOR)
if i==0:
height, width, _ = img.shape
else:
resized_image = cv2.resize(img,(width, height))
cv2.imwrite(f_name,resized_image)
def convert_img(f_name):
img = cv2.imread(f_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)
def merge_sharp1_into2(timg1, timg2, trans1to2, verbose=False):
curr_img = timg2.clone()
warped = KG.homography_warp(timg1, torch.inverse(trans1to2), timg1.shape[-2:])
mask1 = K.filters.laplacian(K.color.rgb_to_grayscale(timg1), 7).abs()
mask1_norm = (mask1-mask1.min()) / (mask1.max() - mask1.min())
mask1_blur = K.filters.gaussian_blur2d(mask1_norm, (9,9), (1.6, 1.6))
mask1_blur = mask1_blur / mask1_blur.max()
warped_mask = KG.homography_warp(mask1_blur.float(), torch.inverse(trans1to2), timg1.shape[-2:])
curr_img = warped_mask * warped + (1-warped_mask) * curr_img
return curr_img
def img_registration(images):
f_names = [f.name for f in images]
resize_images(f_names)
for i, f_name in tqdm(enumerate(f_names)):
if i == 0:
continue
prev_img = convert_img(f_names[i-1]).to(device)
curr_img = convert_img(f_name).to(device)
model = registrator.register(prev_img, curr_img)
models.append(deepcopy(model.detach()))
models_to_final = [torch.eye(3, device=device)[None]]
for m in models[::-1]:
models_to_final.append(m @ models_to_final[-1])
models_to_final = models_to_final[::-1]
base_img = convert_img(f_names[-1])
curr_img = deepcopy(base_img)
_, layers, height, width = curr_img.shape
video_file = 'video.avi'
video = cv2.VideoWriter(video_file, 0, 1, (width,height))
with torch.no_grad():
for i, image in tqdm(enumerate(f_names)):
timg = convert_img(image)
curr_img = merge_sharp1_into2(timg.to(device), curr_img.to(device), models_to_final[i].to(device))
video.write(cv2.cvtColor(K.tensor_to_image(curr_img.float()*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
video.release()
return K.tensor_to_image(curr_img.float()), video_file
title = 'Image Registration with Kornia!'
description = '''Image registration is the process of transforming different sets of data into one coordinate system. Data may be multiple photographs, data from different sensors, times, depths, or viewpoints. It is used in computer vision, medical imaging, and compiling and analyzing images and data from satellites. Registration is necessary in order to be able to compare or integrate the data obtained from these different measurements.
*Note that you can upload only image files, e.g. jpg, png etc and all images should have same width and height!*
Learn more about [image registration and Kornia](https://kornia.readthedocs.io/en/latest/applications/image_registration.html)'''
examples = [["IMG_3020.JPG", "IMG_3027.JPG", "IMG_3034.JPG", "IMG_3040.JPG", "IMG_3058.JPG", "IMG_3070.JPG", "IMG_3083.JPG", "IMG_3100.JPG", "IMG_3106.JPG", "IMG_3112.JPG"]]
iface = gr.Interface(
img_registration,
inputs='files',
outputs=["image", gr.Video()],
allow_flagging="never",
title=title,
description=description
)
if __name__ == "__main__":
iface.launch(show_error=True)
|