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

# Model loading function
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

# Prediction function
def predict(pilimg, conf_thresh, iou_thresh):
    try:
        result = detection_model.predict(pilimg, conf=conf_thresh, iou=iou_thresh)
        img_bgr = result[0].plot()
        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL
        
        # Generate a detection summary
        detections = result[0].boxes.data  # Get detections
        if len(detections) > 0:
            summary = [f"{int(cls)}: {conf:.2f}" for *_, conf, cls in detections]
            summary_str = "\n".join(summary)
        else:
            summary_str = "No detections"
        
        return out_pilimg, summary_str
    except Exception as e:
        return None, f"Error: {str(e)}"

# Repository ID for model
REPO_ID = "ITI107-2024S2/2359747V"
detection_model = load_model(REPO_ID)

# Gradio interface
title = "2359747V Jason Wong YOLOv8s Object Detection App"
description = "Upload an image to detect objects using a YOLOv8s model. Adjust the confidence and IoU thresholds for customisation."

# Custom CSS for a futuristic theme
custom_css = """
body {
    background: linear-gradient(135deg, #1a1a2e, #16213e); /* Dark gradient */
    color: #00f5d4; /* Neon cyan text */
    font-family: 'Arial', sans-serif; /* Modern font */
}

h1, h2, h3, .title {
    color: #00f5d4; /* Neon cyan headers */
    text-shadow: 0px 0px 10px #00f5d4; /* Glow effect */
}

.label {
    color: #f72585; /* Neon pink for labels */
    font-weight: bold;
}

input, button, .gr-button {
    background-color: #1f4068; /* Dark button background */
    color: #e0e1dd; /* Light text for buttons */
    border: 1px solid #00f5d4; /* Neon border */
    border-radius: 5px;
}

.gr-slider > .slider-label {
    color: #00f5d4; /* Neon cyan slider labels */
}

.gr-slider input {
    accent-color: #f72585; /* Neon pink slider accent */
}

.gr-textbox textarea {
    background-color: #1f4068; /* Dark textbox background */
    color: #e0e1dd; /* Light text for textbox */
    border: 1px solid #00f5d4; /* Neon border */
}
"""

interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Slider(0.1, 1.0, 0.5, step=0.05, label="Confidence Threshold"),
        gr.Slider(0.1, 1.0, 0.6, step=0.05, label="IoU Threshold"),
    ],
    outputs=[
        gr.Image(type="pil", label="Detected Image"),
        gr.Textbox(label="Detection Summary, 0=Soldier and 1=Tank"),
    ],
    title=title,
    description=description,
    css=custom_css,  
)

# Launch the interface
interface.launch(share=True)