import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO

model = None

def predict_image(img, conf_threshold, iou_threshold, model_name):
    """Predicts objects in an image using YOLOv8m Defence model with adjustable confidence and IOU thresholds."""
    global model
    model = YOLO(model_name)
    results = model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        show_labels=True,
        show_conf=True,
        imgsz=640,
    )
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
    return im

# Custom CSS for font styling
css = """
body, .gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}
.gr-button {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
    font-weight: 500 !important;
}
.gr-box h1 {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
    font-weight: 600 !important;
}
.gr-box p {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}
"""

iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
        gr.Radio(choices=["yolov8m_defence.pt"], label="Model Name", value="yolov8m_defence.pt"),
    ],
    outputs=gr.Image(type="pil", label="Detection Results"),
    title="YOLOv8m Defence Object Detection",
    description="""
    Upload images to detect military and civilian vehicles, aircraft, and ships using our fine-tuned YOLOv8m model.
    
    **Detectable Objects (18 categories):** Aircraft (cargo, commercial, fighter, helicopter, etc.), 
    Vehicles (car, truck, tank, bus, van), Ships (cargo, yacht, cruise, warship, sailboat), 
    and specialized items (drone, missile).
    
    Developed for DSTA Brainhack 2025 - TIL-AI Category (Semi-Finalist)
    """,
    examples=[
        ["examples/test1.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test2.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test3.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test4.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test5.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test6.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test7.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test8.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test9.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
        ["examples/test10.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
    ],
    css=css
)

iface.launch(share=True)