from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os

# Define the Hugging Face repository ID and the model path
REPO_ID = "aleong888/snoopy"  # Replace with your Hugging Face repository ID

# Function to load the model from the repository
def load_model(repo_id):
    download_dir = snapshot_download(repo_id)
    print("Downloaded model to:", download_dir)
    model_path = os.path.join(download_dir, "best_int8_openvino_model")
    print("Model path:", model_path)
    detection_model = YOLO(model_path, task='detect')
    return detection_model

# Function to make predictions on input images
def predict(pilimg):
    result = detection_model.predict(pilimg, conf=0.5, iou=0.4)
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
    return out_pilimg

# Load the model
detection_model = load_model(REPO_ID)

# Create the Gradio interface
gr.Interface(fn=predict,
             inputs=gr.Image(type="pil", label="Upload an image"),
             outputs=gr.Image(type="pil", label="Detection Result"),
             title="Snoopy Detection Model",
             description="Upload an image to detect Snoopy and Woodstock characters. This application uses a YOLOv8 model trained with OpenVINO optimization."
             ).launch(share=True)