Update app.py

app.py

@@ -6,34 +6,11 @@ import numpy as np
 import matplotlib.pyplot as plt
 from scipy.cluster.vq import kmeans
 
-def get_coordinates_from_mask(mask_in):
-    x_y = np.where(mask_in != [0,0,0,255])[:2]
-    x_y = np.column_stack((x_y[1], x_y[0]))
-    x_y = np.float32(x_y)
-    centroids,_ = kmeans(x_y,4)
-    centroids = np.int64(centroids)
-    return centroids
+# ... (previous functions remain the same)
 
-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])
-    
-    left_coords = c_list[:2]
-    right_coords = c_list[-2:]
-    
-    top_left, bottom_left = get_top_bottom_coordinates(left_coords)
-    top_right, bottom_right = get_top_bottom_coordinates(right_coords)
-    
-    return top_left, top_right, bottom_right, bottom_left
-
-def infer(image_input, dst_height: str, dst_width: str):
-    image_in = image_input["image"]
-    mask_in = image_input["mask"]
+def infer(image_input, mask_input, dst_height: str, dst_width: str):
+    image_in = image_input
+    mask_in = mask_input
     torch_img = K.utils.image_to_tensor(image_in).float() / 255.0
     
     centroids = get_coordinates_from_mask(mask_in)
@@ -78,14 +55,14 @@ description = """In this space you can warp an image using perspective transform
 
 example_mask = np.zeros((327, 600, 4), dtype=np.uint8)
 example_mask[:, :, 3] = 255
-example_image_dict = {"image": "bruce.png", "mask": example_mask}
 
 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")
+        image_input = gr.Image(label="Input Image", type="numpy")
+        mask_input = gr.Image(label="Mask", source="canvas", shape=(600, 327), tool="sketch", interactive=True)
         output_plot = gr.Plot(label="Output")
     
     with gr.Row():
@@ -95,13 +72,13 @@ with gr.Blocks() as demo:
     submit_button = gr.Button("Submit")
     submit_button.click(
         fn=infer,
-        inputs=[image_input, dst_height, dst_width],
+        inputs=[image_input, mask_input, dst_height, dst_width],
         outputs=output_plot
     )
     
     gr.Examples(
-        examples=[[example_image_dict, "64", "128"]],
-        inputs=[image_input, dst_height, dst_width],
+        examples=[["bruce.png", example_mask, "64", "128"]],
+        inputs=[image_input, mask_input, dst_height, dst_width],
         outputs=output_plot,
         fn=infer,
         cache_examples=True