Update app.py

app.py

@@ -158,63 +158,59 @@ iface = gr.Interface(
 
 # Launch the Gradio app
 iface.launch()"""
-
 from PIL import Image
 import numpy as np
-import gradio as gr
 from ultralytics import YOLO
+import gradio as gr
 
 # Load the YOLO model
 MODEL_URL = 'https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt'
 model = YOLO(MODEL_URL)
 
-# Define the prediction function
+# Define the prediction function for Gradio
 def predict(image):
     try:
-        # 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.")
+        # Convert PIL image to NumPy array
+        image_array = np.array(image)
 
         # Perform prediction
-        results = model.predict(image)  # Perform object detection
-        print("YOLO prediction completed.")
+        results = model(image_array)
 
-        # Print detected class names
-        detected_classes = results.names  # Get class names for each detected object
-        print("Detected classes:", detected_classes)
+        # Access the first result
+        result = results[0]
 
-        # Print the names of the detected classes
-        for i, name in enumerate(detected_classes):
-            print(f"Detected object {i+1}: {name}")
+        # Extract detected classes
+        detected_classes = [model.names[int(cls)] for cls in result.boxes.cls]
+        print(f"Detected classes: {detected_classes}")
 
         # Render bounding boxes on the image
-        results.render()
-        print("Bounding boxes rendered.")
+        annotated_image = result.plot()
 
-        # Convert back to PIL for output
-        output_image = Image.fromarray(results.imgs[0])
-        print("Output image type:", type(output_image))
+        # Convert the annotated image to PIL format
+        output_image = Image.fromarray(annotated_image)
 
-        return output_image
+        # Return the annotated image and detected classes as output
+        return output_image, detected_classes
 
     except Exception as e:
         print("Error during prediction:", str(e))
-        return None
+        return None, ["Error during processing"]
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an Image"),  # Upload input as PIL Image
-    outputs=gr.Image(type="pil", label="Predicted Image with Bounding Boxes"),  # Output image
-    title="Object Detection App",
-    description="Upload an image, and the YOLO model will detect objects in it."
+    inputs=gr.Image(type="pil", label="Upload an Image"),  # Input image as PIL
+    outputs=[
+        gr.Image(type="pil", label="Predicted Image with Bounding Boxes"),  # Annotated image
+        gr.Label(label="Detected Classes")  # Detected classes
+    ],
+    title="YOLO Object Detection App",
+    description="Upload an image, and the YOLO model will detect objects and annotate them with bounding boxes and class labels."
 )
 
 # Launch the Gradio app
 iface.launch()
 
 
+
+