File size: 1,871 Bytes
fa53d0e |
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 |
import gradio as gr
import kornia as K
import numpy as np
import kornia.geometry as KG
import torch
import cv2
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]
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)
transformed: torch.Tensor = KG.transform.warp_affine(image_batch, affine_matrix[:, :2], dsize=image_batch.shape[2:])
np_images: list = [K.tensor_to_image(f*255).astype(np.uint8) for f in transformed]
final_images: np.ndarray = cv2.hconcat(np_images)
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() |