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

REPO_ID = "jiaxinnnnn/Interior-Style-Classification"
WEIGHTS_FILENAME = "Best_Accuracy.pt"

def load_model(repo_id: str) -> YOLO:
    """
    Downloads the model repo from Hugging Face Hub into the Space cache
    and loads the YOLO model from the weights file.
    """
    download_dir = snapshot_download(repo_id=repo_id)
    weights_path = os.path.join(download_dir, WEIGHTS_FILENAME)

    if not os.path.exists(weights_path):
        raise FileNotFoundError(
            f"Cannot find {WEIGHTS_FILENAME} inside downloaded repo: {download_dir}\n"
            f"Files found: {os.listdir(download_dir)}"
        )

    return YOLO(weights_path, task="classify")

classification_model = load_model(REPO_ID)

def predict(pilimg: Image.Image):
    results = classification_model.predict(pilimg)

    probs = results[0].probs
    class_id = probs.top1
    confidence = probs.top1conf.item()
    class_name = classification_model.names[class_id]

    return {
        "style": class_name.title(),
        "confidence": confidence
    }


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an interior image"),
    outputs=gr.JSON(label="Prediction"),
    title="Interior Style Classification (YOLOv8)",
    description="Upload an image and the model will classify the interior design style."
)

demo.launch()