Update app.py

app.py

@@ -1,22 +1,37 @@
 import gradio as gr
+from PIL import Image
+import numpy as np
 from TheDistanceAssessor import run, load_segformer, load_yolo
 
 def process_image(input_image):
-    image_size = [1024,1024]
-    target_distances = [650,1000,2000]
+    image_size = [1024, 1024]
+    target_distances = [650, 1000, 2000]
     num_ys = 10
 
     PATH_model_seg = 'SegFormer_B3_1024_finetuned.pth'
     PATH_model_det = 'yolov8s.pt'
     model_seg = load_segformer(PATH_model_seg)
     model_det = load_yolo(PATH_model_det)
-    output_image = run(input_image, model_seg, model_det, image_size, target_distances, num_ys = num_ys)
+
+    # Convert the input image to a numpy array if necessary
+    if isinstance(input_image, str):
+        # If the input is a file path, open the image using PIL
+        input_image = Image.open(input_image)
+    
+    if isinstance(input_image, Image.Image):
+        # Convert PIL image to numpy array
+        input_image = np.array(input_image)
+    
+    if not isinstance(input_image, np.ndarray):
+        raise ValueError("Input image is not a valid format (string path, PIL image, or numpy array)")
+
+    output_image = run(input_image, model_seg, model_det, image_size, target_distances, num_ys=num_ys)
     return output_image
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=process_image,  # The function to be called
-    inputs=gr.Image(type="pil"),  # Input type
+    inputs=gr.Image(type="auto"),  # Input type auto-detects the type
     outputs=gr.Image(type="numpy"),  # Output type
     title="Image Processor",  # Title of the interface
     description="Upload an image and get a processed image as output."  # Description of the interface
@@ -24,4 +39,4 @@ iface = gr.Interface(
 
 # Launch the interface
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()