Update app.py

app.py

@@ -5,6 +5,7 @@ import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 from scipy.cluster.vq import kmeans
+from PIL import Image
 
 def get_coordinates_from_mask(mask_in):
     x_y = np.where(mask_in != [0,0,0,255])[:2]
@@ -32,8 +33,13 @@ 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):
-    image_in = image_input["image"]
-    mask_in = image_input["mask"]
+    if isinstance(image_input, dict):
+        image_in = image_input["image"]
+        mask_in = image_input["mask"]
+    else:
+        image_in = image_input
+        mask_in = np.zeros_like(image_in)
+    
     torch_img = K.utils.image_to_tensor(image_in).float() / 255.0
     
     centroids = get_coordinates_from_mask(mask_in)
@@ -76,16 +82,15 @@ description = """In this space you can warp an image using perspective transform
 4. Click Submit to run the demo
 """
 
-example_mask = np.zeros((327, 600, 4), dtype=np.uint8)
-example_mask[:, :, 3] = 255
-example_image_dict = {"image": "bruce.png", "mask": example_mask}
+# Load the example image
+example_image = Image.open("bruce.png")
+example_image_np = np.array(example_image)
 
 inputs = [
-    gr.Image(type="numpy", label="Input Image"),
+    gr.Image(type="numpy", label="Input Image", tool="sketch"),
     gr.Textbox(label="Destination Height", value="64"),
     gr.Textbox(label="Destination Width", value="128"),
 ]
-
 outputs = gr.Plot(label="Output")
 
 gr.Interface(
@@ -94,6 +99,6 @@ gr.Interface(
     outputs=outputs,
     title="Homography Warping",
     description=description,
-    examples=[[example_image_dict, "64", "128"]],
+    examples=[[example_image_np, "64", "128"]],
     cache_examples=True
 ).launch()