| import gradio as gr | |
| import kornia as K | |
| import numpy as np | |
| import kornia.geometry as KG | |
| import torch | |
| def geometry_transform(images: list, | |
| translation: float, | |
| scale: float, | |
| angle: float) -> np.ndarray: | |
| file_names: list = [f.name for f in images] | |
| image_list: list = [K.io.load_image(f, K.io.ImageLoadType(0)).float().unsqueeze(0)/255 for f in file_names] | |
| if len(image_list) > 1: | |
| image_list = [K.geometry.resize(x, x.shape[-2:], antialias=True) for x in image_list] | |
| image_batch: torch.Tensor = torch.cat(image_list, 0) | |
| center: torch.Tensor = torch.tensor([x.shape[1:] for x in image_batch])/2 | |
| translation = torch.tensor(translation).repeat(len(image_list), 2) | |
| scale = torch.tensor(scale).repeat(len(image_list), 2) | |
| angle = torch.tensor(angle).repeat(len(image_list)) | |
| affine_matrix: torch.Tensor = KG.get_affine_matrix2d(translation, center, scale, angle) | |
| with torch.inference_mode(): | |
| transformed: torch.Tensor = KG.transform.warp_affine(image_batch, affine_matrix[:, :2], dsize=image_batch.shape[2:]) | |
| concat_images: list = torch.cat([x for x in transformed], dim=-1) | |
| final_images: np.ndarray = K.tensor_to_image(concat_images*255).astype(np.uint8) | |
| return final_images | |
| def main(): | |
| title = """ | |
| <h1 align="center"> | |
| Geometry Image Transforms with Kornia! | |
| </h1> | |
| """ | |
| with gr.Blocks() as demo: | |
| gr.Markdown(title) | |
| with gr.Row(): | |
| images_input = gr.Files() | |
| with gr.Column(): | |
| translation = gr.Number(label= "Translation") | |
| scale = gr.Number(label = "Scale", value= 1.0) | |
| angle = gr.Number(label = "Angle") | |
| button = gr.Button('Transform') | |
| result = gr.Image() | |
| button.click( | |
| geometry_transform, | |
| inputs=[images_input,translation, scale, angle], | |
| outputs=result | |
| ) | |
| demo.launch() | |
| if __name__ == '__main__': | |
| main() |