Update app.py

app.py

@@ -1,42 +1,38 @@
 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, conf_threshold, iou_threshold):
-    """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds."""
+def predict_image(img):
+    """Classifies an image using a YOLOv8 model."""
     results = model.predict(
         source=img,
-        conf=conf_threshold,
-        iou=iou_threshold,
-        show_labels=True,
-        show_conf=True,
         imgsz=640,
+        conf=0.25,  # You can adjust this confidence threshold
     )
-
-    for r in results:
-        im_array = r.plot()
-        im = Image.fromarray(im_array[..., ::-1])
-
-    return im
-
+    
+    # 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"
 
 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.Image(type="pil", label="Result"),
+    outputs=gr.Text(label="Classification Result"),
     title="GP Wolf Classifier",
-    description="Upload images for inference.",
+    description="Upload images for classification.",
     examples=[
-        [ASSETS / "gp.jpg", 0.25, 0.45],
-        [ASSETS / "wolf.jpg", 0.25, 0.45],
+        [ASSETS / "https://huggingface.co/spaces/gpbhupinder/test/blob/main/gp.jpg"],
+        [ASSETS / "https://huggingface.co/spaces/gpbhupinder/test/blob/main/wolf.jpg"],
     ],
 )