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."""
    results = model.predict(
        source=img,
        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"

iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
    ],
    outputs=gr.Text(label="Classification Result"),
    title="GP Wolf Classifier",
    description="Upload images for classification.",
    examples=[
         ["gp.jpg"],
        ["wolf.jpg"],
    ],
)

if __name__ == "__main__":
    iface.launch()