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_dir = snapshot_download(repo_id)
    print(download_dir)
    path = os.path.join(download_dir, "best_int8_openvino_model")
    print(path)
    detection_model = YOLO(path, task='detect')
    return detection_model


def predict(pilimg, confidence, iou):
    source = pilimg
    result = detection_model.predict(source, conf=confidence, iou=iou)
    img_bgr = result[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
    return out_pilimg


REPO_ID = "skngew/9053220B"
detection_model = load_model(REPO_ID)

# Student ID
student_id = "Student ID: 9053220B"

# Create the Gradio interface
def create_interface():
    # Persistent state for default values
    confidence_default = gr.State(0.5)
    iou_default = gr.State(0.6)

    interface = gr.Interface(
        fn=predict,
        inputs=[
            gr.Image(type="pil", label="Input Image"),
            gr.Slider(0, 1, value=confidence_default.value, label="Confidence Threshold"),  # Default to 0.5
            gr.Slider(0, 1, value=iou_default.value, label="IOU Threshold")  # Default to 0.6
        ],
        outputs=gr.Image(type="pil", label="Output Image"),
        title="Object Detection with YOLOv8",
        description=student_id,
        live=False,
    )

    return interface


# Launch the Gradio app
app_interface = create_interface()
app_interface.launch(share=True)