Update app.py

app.py

@@ -1,37 +1,40 @@
 import gradio as gr
 import PIL.Image as Image
+
 from ultralytics import ASSETS, YOLO
 
 model = YOLO("https://huggingface.co/spaces/gpbhupinder/test/blob/main/model_-%2023%20june%202024%2019_22.pt")
 
+
 def predict_image(img):
-    """Classifies an image using a YOLOv8 model."""
+    """Predicts objects in an image using a YOLOv8 model."""
     results = model.predict(
         source=img,
+        
+        show_labels=True,
+        show_conf=True,
         imgsz=640,
-        conf=0.25,  # You can adjust this confidence threshold
     )
-    
-    # Get the top prediction
-    if results and len(results[0].boxes) > 0:
-        top_prediction = results[0].boxes[0]
-        class_id = int(top_prediction.cls)
-        confidence = float(top_prediction.conf)
-        class_name = model.names[class_id]
-        return f"{class_name} (Confidence: {confidence:.2f})"
-    else:
-        return "No classification made"
+
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
+
+    return im
+
 
 iface = gr.Interface(
     fn=predict_image,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
+        # gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
+        # gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
     ],
-    outputs=gr.Text(label="Classification Result"),
+    outputs=gr.Image(type="pil", label="Result"),
     title="GP Wolf Classifier",
-    description="Upload images for classification.",
+    description="Upload images for inference.",
     examples=[
-         ["gp.jpg"],
+        ["gp.jpg"],
         ["wolf.jpg"],
     ],
 )