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

model_path = "best_int8_openvino_model"

# Function to load the model from Hugging Face repo
def load_model(repo_id):
    download_dir = snapshot_download(repo_id)
    print(download_dir)
    path  = os.path.join(download_dir, model_path)
    print(path)
    detection_model = YOLO(path, task='detect')
    return detection_model

# Prediction function with dynamic conf and iou
def predict(pilimg, conf, iou):
    source = pilimg
    # Run object detection with user-specified conf and iou
    result = detection_model.predict(source, conf=conf, iou=iou)
    
    # Visualize the result with bounding boxes
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
    
    return out_pilimg

# Hugging Face repo ID
REPO_ID = "Sathyadithyarithi/RRR_detection"
detection_model = load_model(REPO_ID)

# Gradio Interface

gr.Markdown("### Set Confidence (conf) and Intersection over Union (IoU) Thresholds")

gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Image to Detect Objects"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence (conf)", info="Set the confidence threshold for object detection."),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.05, label="Intersection over Union (IoU)", info="Set the IoU threshold to filter detections.")
    ],
    outputs=gr.Image(type="pil", label="Detected Image"),
    title="Object Detection App for RRR Charan and NTR"
   ).launch(share=True)