| 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() |
