"""
Gradio App for PPE Compliance Detection
Run: gradio gradio_app.py
"""
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
from PIL import Image
import cv2
import numpy as np

MODEL_ID = "baskarmother/yolov8-ppe-construction"
CLASS_NAMES = [
    'barricade', 'dumpster', 'excavators', 'gloves', 'hardhat', 'mask',
    'no-hardhat', 'no-mask', 'no-safety vest', 'person', 'safety net',
    'safety shoes', 'safety vest', 'dump truck', 'mini-van', 'truck', 'wheel loader'
]

def get_color(label):
    if label in ["hardhat", "mask", "safety vest", "gloves", "safety shoes"]:
        return (0, 255, 0)  # Green = compliant PPE
    elif label in ["no-hardhat", "no-mask", "no-safety vest"]:
        return (0, 0, 255)  # Red = violation
    else:
        return (255, 0, 0)  # Blue = other

# Load model once
print("Loading model from Hub...")
weights_path = hf_hub_download(MODEL_ID, "best.pt")
model = YOLO(weights_path)
print("Model loaded!")

def detect_ppe(image, conf_threshold=0.25):
    results = model(image, conf=conf_threshold)
    result = results[0]
    
    # Convert PIL to numpy array
    img = np.array(image)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    detections = []
    compliance = {"hardhat": "NOT DETECTED", "mask": "NOT DETECTED", "safety_vest": "NOT DETECTED"}
    detected_classes = set()
    
    for box in result.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cls_id = int(box.cls[0])
        conf = float(box.conf[0])
        label = CLASS_NAMES[cls_id]
        detected_classes.add(label)
        
        color = get_color(label)
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
        text = f"{label} {conf:.2f}"
        cv2.putText(img, text, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        detections.append(f"{label}: {conf:.2f}")
    
    # Compliance check
    if "hardhat" in detected_classes:
        compliance["hardhat"] = "✅ COMPLIANT"
    if "no-hardhat" in detected_classes:
        compliance["hardhat"] = "❌ VIOLATION"
    if "mask" in detected_classes:
        compliance["mask"] = "✅ COMPLIANT"
    if "no-mask" in detected_classes:
        compliance["mask"] = "❌ VIOLATION"
    if "safety vest" in detected_classes:
        compliance["safety_vest"] = "✅ COMPLIANT"
    if "no-safety vest" in detected_classes:
        compliance["safety_vest"] = "❌ VIOLATION"
    
    # Convert back to RGB
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    output_image = Image.fromarray(img_rgb)
    
    report = f"""### PPE Compliance Report
- **Hard Hat**: {compliance['hardhat']}
- **Mask**: {compliance['mask']}
- **Safety Vest**: {compliance['safety_vest']}

### Detections ({len(detections)} objects)
{chr(10).join(['- ' + d for d in detections]) if detections else 'No objects detected'}
"""
    return output_image, report

with gr.Blocks(title="PPE Compliance Detection") as demo:
    gr.Markdown("# 🦺 PPE Compliance Detection for Construction Sites")
    gr.Markdown("Upload a construction site image to detect Personal Protective Equipment (PPE) compliance.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Image")
            conf_slider = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="Confidence Threshold")
            detect_btn = gr.Button("Detect PPE", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Detected Objects")
            output_text = gr.Markdown(label="Compliance Report")
    
    detect_btn.click(fn=detect_ppe, inputs=[input_image, conf_slider], outputs=[output_image, output_text])
    
    gr.Markdown("""
    ## Classes Detected
    - ✅ **Compliant PPE**: hardhat, mask, safety vest, gloves, safety shoes
    - ❌ **Violations**: no-hardhat, no-mask, no-safety vest
    - 🔵 **Other**: person, equipment, vehicles, barricades
    """)

if __name__ == "__main__":
    demo.launch()