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

# Load model once at startup
print("Loading YOLOv8m Defence model...")
model = YOLO("yolov8m_defence.pt")

# Set device and optimize for CPU inference
if torch.cuda.is_available():
    device = 'cuda'
    print("Using GPU acceleration")
else:
    device = 'cpu'
    print("Using CPU inference")
    torch.set_num_threads(2)

model.to(device)

def predict_image(img, conf_threshold, iou_threshold):
    """Predicts objects in an image using YOLOv8m Defence model with adjustable confidence and IOU thresholds."""
    try:
        results = model.predict(
            source=img,
            conf=conf_threshold,
            iou=iou_threshold,
            show_labels=True,
            show_conf=True,
            imgsz=640,
            verbose=False,
            device=device
        )
        
        for r in results:
            im_array = r.plot()
            im = Image.fromarray(im_array[..., ::-1])
        
        return im
    
    except Exception as e:
        print(f"Error during prediction: {e}")
        return img

# 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;
}
"""

# Create interface
iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence threshold"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="IoU threshold"),
    ],
    outputs=gr.Image(type="pil", label="Detection Results"),
    title="Defence Object Detection",
    description="""
    Upload images to detect military and civilian vehicles, aircraft, and ships using our fine-tuned YOLOv8m model.
    
    **Model Card:** [spencercdz/YOLOv8m_defence](https://huggingface.co/spencercdz/YOLOv8m_defence)
    
    **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).
    
    **Performance:** Running on Free Tier - inference may take up to 40 seconds per image.
    """,
    examples=[
        ["examples/test1.jpg", 0.4, 0.5],
        ["examples/test2.jpg", 0.4, 0.5],
        ["examples/test3.jpg", 0.4, 0.5],
        ["examples/test4.jpg", 0.4, 0.5],
        ["examples/test5.jpg", 0.4, 0.5],
        ["examples/test6.jpg", 0.4, 0.5],
        ["examples/test7.jpg", 0.4, 0.5],
        ["examples/test8.jpg", 0.4, 0.5],
        ["examples/test9.jpg", 0.4, 0.5],
        ["examples/test10.jpg", 0.4, 0.5],
    ],
    css=css,
    cache_examples=True,
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch(share=True)