Update app.py

app.py

@@ -1,57 +1,44 @@
+import gradio as gr
+from ultralytics import YOLO
 from PIL import Image
 import numpy as np
-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_URL = 'best.pt'
+MODEL_URL = "https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt"
 model = YOLO(MODEL_URL)
 
-# Define the prediction function for Gradio
-def predict(image):
+# Define the prediction function
+def predict(input_img):
     try:
-        # Convert PIL image to NumPy array
-        image_array = np.array(image)
+        # Convert PIL Image to NumPy array
+        image_array = np.array(input_img)
 
-        # Perform prediction
+        # Perform inference
         results = model(image_array)
 
-        # Access the first result
-        result = results[0]
-
-        # 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
-        annotated_image = result.plot()
+        # Extract detected class names
+        detected_classes = [model.names[int(cls)] for cls in results[0].boxes.cls]
 
-        # Convert the annotated image to PIL format
-        output_image = Image.fromarray(annotated_image)
-
-        # Return the annotated image and detected classes as output
-        return output_image, detected_classes
+        # Render results on the image
+        results[0].plot()  # Render bounding boxes on the image
+        output_image = Image.fromarray(results[0].orig_img)
 
+        return output_image, {cls: 1.0 for cls in detected_classes}  # Dummy scores for simplicity
     except Exception as e:
-        print("Error during prediction:", str(e))
-        return None, ["Error during processing"]
+        print(f"Error during processing: {e}")
+        return None, {"Error": str(e)}
 
-# Create the Gradio interface
-iface = gr.Interface(
+# Gradio app configuration
+gradio_app = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an Image"),  # Input image as PIL
+    inputs=gr.Image(label="Upload an Image", type="pil"),
     outputs=[
-        gr.Image(type="pil", label="Predicted Image with Bounding Boxes"),  # Annotated image
-        gr.Label(label="Detected Classes")  # Detected classes
+        gr.Image(label="Predicted Image with Bounding Boxes"),  # Rendered image with bounding boxes
+        gr.Label(label="Detected Classes"),  # Detected class names
     ],
     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."
+    description="Upload an image, and the YOLO model will detect objects in it.",
 )
 
-# Launch the Gradio app
-iface.launch()
-
-
-
-
+if __name__ == "__main__":
+    gradio_app.launch()