Update app.py

app.py

@@ -160,6 +160,7 @@ iface = gr.Interface(
 iface.launch()"""
 
 from PIL import Image
+import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 
@@ -170,19 +171,30 @@ model = YOLO(MODEL_URL)
 # Define the prediction function
 def predict(image):
     try:
-        print("Received image:", type(image))  # Check the type of the received image (should be PIL)
-        results = model(image)  # Perform object detection using YOLO
-        
-        # Print details of the results object
-        print("Results:", results)  # Check the contents of results
-        
-        results.render()  # Render bounding boxes on the image
-        
-        output_image = Image.fromarray(results.imgs[0])  # Convert to PIL image
-        print("Predicted image type:", type(output_image))  # Ensure output is PIL Image
+        # Debugging: Check image type
+        print("Input image type:", type(image))
+
+        # Ensure input is converted to a format the YOLO model supports
+        if isinstance(image, Image.Image):  # If PIL image, convert to NumPy
+            image = np.array(image)  # Convert to NumPy array
+            print("Converted image to NumPy array.")
+
+        # Perform prediction
+        results = model.predict(image)  # Perform object detection
+        print("YOLO prediction completed.")
+
+        # Render bounding boxes on the image
+        results.render()
+        print("Bounding boxes rendered.")
+
+        # Convert back to PIL for output
+        output_image = Image.fromarray(results.imgs[0])
+        print("Output image type:", type(output_image))
+
         return output_image
+
     except Exception as e:
-        print("Error during prediction:", e)
+        print("Error during prediction:", str(e))
         return None
 
 # Create the Gradio interface