import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import torch
import numpy as np
from PIL import Image
import os  # Importiere das 'os'-Modul

# Definiere die Model-ID und den Dateinamen
repo_id = "Web4/LS-W4-SX-YOLO-Light-V1-Candy"
model_file = "my_model.pt"

# Hole das Hugging Face Token aus den Umgebungsvariablen (dem Secret)
# Das Token wird nur bei Gated-Repos benötigt.
token = os.environ.get("HF_TOKEN")

# Lade die Modelldatei in den lokalen Cache des Spaces,
# indem du das Token für die Authentifizierung übergibst.
try:
    model_path = hf_hub_download(repo_id=repo_id, filename=model_file, token=token)
    print(f"Model downloaded to: {model_path}")
except Exception as e:
    print(f"Error downloading model: {e}")
    raise

# Lade das Modell von dem lokalen Pfad, den hf_hub_download bereitstellt
model = YOLO(model_path)

# Der Rest des Codes bleibt unverändert
def predict_image(input_image):
    if input_image is None:
        return None
    
    results = model.predict(input_image)
    
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
        return im

demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(label="Upload an image"),
    outputs=gr.Image(label="Object Detection Result"),
    title="YOLO-Light-V1-Candy Demo",
    description="Upload an image to see the object detection model in action."
)

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