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

def load_model(repo_id):
    # Download the snapshot
    download_dir = snapshot_download(repo_id)
    print(f"Downloaded snapshot directory: {download_dir}")

    # Check the contents of the directory
    for root, dirs, files in os.walk(download_dir):
        print(f"Found files: {files}")
    
    # Check for the correct model file and path
    # Ensure the path points to the correct location where the model is saved
    model_path = os.path.join(download_dir, "best_int8_openvino_model")
    
    # Check if the model file exists at the expected location
    if not os.path.exists(model_path):
        print(f"Model file not found at {model_path}")
        return None

    # Load the model using YOLO
    detection_model = YOLO(model_path, task='detect')
    return detection_model


def predict(pilimg):
    result = detection_model.predict(pilimg, conf=0.5, iou=0.6)
    img_bgr = result[0].plot()  # Get image with predictions
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
    return out_pilimg


# Set the repo ID
REPO_ID = "Lesterchia174/Monkey_Durian"

# Load the model
detection_model = load_model(REPO_ID)

if detection_model:
    gr.Interface(fn=predict,
                 inputs=gr.Image(type="pil"),
                 outputs=gr.Image(type="pil")
                 ).launch(share=True)
else:
    print("Model loading failed. Check the model path or file structure.")