Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+"""import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
 import numpy as np
@@ -40,5 +40,51 @@ gradio_app = gr.Interface(
     description="Upload an image, and the YOLO model will detect objects in it.",
 )
 
+if __name__ == "__main__":
+    gradio_app.launch()"""
+
+import gradio as gr
+from ultralytics import YOLO
+from PIL import Image
+import numpy as np
+
+# Load the YOLO model
+MODEL_URL = "https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt"
+model = YOLO(MODEL_URL)
+
+# Define the prediction function
+def predict(input_img):
+    try:
+        # Convert PIL Image to NumPy array
+        image_array = np.array(input_img)
+
+        # Perform inference
+        results = model(image_array)
+
+        # Extract detected class names
+        detected_classes = [model.names[int(cls)] for cls in results[0].boxes.cls]
+
+        # Render results on the image
+        rendered_image = results[0].plot()  # This method returns the image with bounding boxes
+        output_image = Image.fromarray(rendered_image)  # Convert the rendered image to a PIL Image
+
+        return output_image, {cls: 1.0 for cls in detected_classes}  # Dummy scores for simplicity
+    except Exception as e:
+        print(f"Error during processing: {e}")
+        return None, {"Error": str(e)}
+
+# Gradio app configuration
+gradio_app = gr.Interface(
+    predict,
+    inputs=gr.Image(label="Upload an Image", type="pil"),
+    outputs=[
+        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 in it.",
+)
+
 if __name__ == "__main__":
     gradio_app.launch()
+