Update app.py

app.py

@@ -1,13 +1,10 @@
 import gradio as gr
 import torch
 import kornia as K
-
 import cv2
 import numpy as np
-import matplotlib
 import matplotlib.pyplot as plt
-matplotlib.use('Agg')
-from scipy.cluster.vq import kmeans,vq,whiten
+from scipy.cluster.vq import kmeans
 
 def get_coordinates_from_mask(mask_in):
     x_y = np.where(mask_in != [0,0,0,255])[:2]
@@ -15,16 +12,13 @@ def get_coordinates_from_mask(mask_in):
     x_y = np.float32(x_y)
     centroids,_ = kmeans(x_y,4)
     centroids = np.int64(centroids)
-
     return centroids
 
 def get_top_bottom_coordinates(coords):
     top_coord = min(coords, key=lambda x : x[1])
     bottom_coord = max(coords, key=lambda x : x[1])
-    
     return top_coord, bottom_coord
 
-
 def sort_centroids_clockwise(centroids: np.ndarray):    
     c_list = centroids.tolist()
     c_list.sort(key = lambda y : y[0])
@@ -37,75 +31,81 @@ def sort_centroids_clockwise(centroids: np.ndarray):
     
     return top_left, top_right, bottom_right, bottom_left
 
-
-def infer(image_input, dst_height:str, dst_width:str):
+def infer(image_input, dst_height: str, dst_width: str):
     image_in = image_input["image"]
     mask_in = image_input["mask"]
-    torch_img = K.image_to_tensor(image_in)
+    torch_img = K.utils.image_to_tensor(image_in).float() / 255.0
     
     centroids = get_coordinates_from_mask(mask_in)
     ordered_src_coords = sort_centroids_clockwise(centroids)
     # the source points are the region to crop corners
     points_src = torch.tensor([list(ordered_src_coords)], dtype=torch.float32)
-
     # the destination points are the image vertexes
     h, w = int(dst_height), int(dst_width)  # destination size
     points_dst = torch.tensor([[
         [0., 0.], [w - 1., 0.], [w - 1., h - 1.], [0., h - 1.],
     ]], dtype=torch.float32)
-
     # compute perspective transform
-    M: torch.tensor = K.geometry.get_perspective_transform(points_src, points_dst)
-
+    M: torch.tensor = K.geometry.transform.get_perspective_transform(points_src, points_dst)
     # warp the original image by the found transform
     torch_img = torch.stack([torch_img],)
-    img_warp: torch.tensor = K.geometry.warp_perspective(torch_img.float(), M, dsize=(h, w))
-
+    img_warp: torch.tensor = K.geometry.transform.warp_perspective(torch_img, M, dsize=(h, w))
     
     # convert back to numpy
-    img_np = K.tensor_to_image(torch_img.byte())
-    img_warp_np: np.ndarray = K.tensor_to_image(img_warp.byte())
-
+    img_np = K.utils.tensor_to_image(torch_img[0])
+    img_warp_np: np.ndarray = K.utils.tensor_to_image(img_warp[0])
     # draw points into original image
     for i in range(4):
         center = tuple(points_src[0, i].long().numpy())
         img_np = cv2.circle(img_np.copy(), center, 5, (0, 255, 0), -1)
-
     # create the plot
     fig, axs = plt.subplots(1, 2, figsize=(16, 10))
     axs = axs.ravel()
-
     axs[0].axis('off')
     axs[0].set_title('image source')
     axs[0].imshow(img_np)
-
     axs[1].axis('off')
     axs[1].set_title('image destination')
     axs[1].imshow(img_warp_np)
-
     return fig
 
-
 description = """In this space you can warp an image using perspective transform with the Kornia library as seen in [this tutorial](https://kornia.github.io/tutorials/#category=Homography).
-
-1. Upload an image (e.g. [this one](https://github.com/kornia/data/raw/main/bruce.png))
+1. Upload an image or use the example provided
 2. Set 4 points into the image with your cursor, which define the area to warp
 3. Set a desired output size (or go with the default)
 4. Click Submit to run the demo
 """
 
-example_mask = np.empty((327,600,4))
-example_mask[:] = [0,0,0,255]
+example_mask = np.zeros((327, 600, 4), dtype=np.uint8)
+example_mask[:, :, 3] = 255
 example_image_dict = {"image": "bruce.png", "mask": example_mask}
 
-Iface = gr.Interface(
-    fn=infer,
-    inputs=[gr.components.Image(tool="sketch"),
-            gr.components.Textbox(label="Destination Height", value="64"),
-            gr.components.Textbox(label="Destination Width", value="128"),
-            ],
-    outputs=gr.components.Plot(),
-    #examples=[["bruce.png", example_mask, "64", "128"]],
-    title="Homography Warping",
-    description=description,
-).launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# Homography Warping")
+    gr.Markdown(description)
+    
+    with gr.Row():
+        image_input = gr.Image(tool="sketch", type="numpy", label="Input Image")
+        output_plot = gr.Plot(label="Output")
+    
+    with gr.Row():
+        dst_height = gr.Textbox(label="Destination Height", value="64")
+        dst_width = gr.Textbox(label="Destination Width", value="128")
+    
+    submit_button = gr.Button("Submit")
+    submit_button.click(
+        fn=infer,
+        inputs=[image_input, dst_height, dst_width],
+        outputs=output_plot
+    )
+    
+    gr.Examples(
+        examples=[[example_image_dict, "64", "128"]],
+        inputs=[image_input, dst_height, dst_width],
+        outputs=output_plot,
+        fn=infer,
+        cache_examples=True
+    )
+
+if __name__ == "__main__":
+    demo.launch()