import os
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# Download private model using HF token from secret
hf_token = os.environ.get("HF_TOKEN")
model_path = hf_hub_download(
    repo_id="Lookingsoft-team/object_detection", filename="yolov8n.pt", token=hf_token
)

# Load model
model = YOLO(model_path)


def detect_objects(image):
    results = model(image)
    return results[0].plot()


demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title="Object Detection",
    description="Upload an image to detect objects.",
)

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