Create app.py

app.py

@@ -0,0 +1,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 check_image_sizes(f_names):
+ sizes = []
+ for f_name in f_names:
+ img = cv2.imread(f_name, cv2.IMREAD_COLOR)
+ sizes.append(img.shape)
+
+ return len(set(sizes))==1
+
+
+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]
+ # print(check_image_sizes(f_names))
+ # if not check_image_sizes(f_names):
+ # return "All images should be same size!"
+
+ 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 registraion 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)