import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import spaces

# Load YOLO model
model = YOLO("best.pt")

@spaces.GPU
def detect_ppe(image):
    """
    Detect PPE violations in the image
    """
    if image is None:
        return None
    
    # Run YOLO inference
    results = model(image, conf=0.4)
    
    # Get the annotated image
    annotated_image = results[0].plot()
    
    # Convert BGR to RGB for Gradio
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    return annotated_image

# Custom CSS for professional look
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Cairo:wght@400;500;600;700&display=swap');

* {
    font-family: 'Cairo', sans-serif !important;
}

.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%) !important;
}

#component-0 {
    background: white !important;
    border-radius: 16px !important;
    padding: 2rem !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1) !important;
}

.gr-button-primary {
    background: linear-gradient(135deg, #267649, #339966) !important;
    border: none !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    padding: 12px 32px !important;
    border-radius: 8px !important;
    transition: all 0.3s !important;
}

.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 12px rgba(38, 118, 73, 0.3) !important;
}

h1 {
    color: #267649 !important;
    font-size: 32px !important;
    font-weight: 700 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
}

.gr-prose p {
    color: #555 !important;
    font-size: 16px !important;
    text-align: center !important;
}

.gr-box {
    border-radius: 12px !important;
    border: 2px solid #e0e0e0 !important;
}

.gr-input-label {
    color: #267649 !important;
    font-weight: 600 !important;
    font-size: 16px !important;
}

footer {
    display: none !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="green")) as demo:
    gr.HTML("""
        <div style="text-align: center; padding: 20px 0;">
            <h1>🛡️ نظام امتثال لكشف معدات السلامة</h1>
            <p style="font-size: 18px; color: #666; margin-top: 10px;">
                اختبر نظام الذكاء الاصطناعي لكشف عدم ارتداء القفازات، الكمامات، والقبعات
            </p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                sources=["webcam", "upload"], 
                type="numpy",
                label="📷 الكاميرا أو رفع صورة",
                height=400
            )
            submit_btn = gr.Button("🔍 ابدأ الكشف", variant="primary", size="lg")
        
        with gr.Column():
            output_image = gr.Image(
                label="✅ نتائج الكشف",
                type="numpy",
                height=400
            )
    
    gr.HTML("""
        <div style="text-align: center; padding: 20px; background: #f5f5f5; border-radius: 12px; margin-top: 20px;">
            <p style="margin: 0; color: #666; font-size: 14px;">
                ✅ اسمح بالوصول للكاميرا عند الطلب • 
                👤 ضع نفسك أمام الكاميرا • 
                🎯 سيقوم النظام بكشف المخالفات تلقائياً
            </p>
        </div>
    """)
    
    # Event handlers
    submit_btn.click(
        fn=detect_ppe,
        inputs=input_image,
        outputs=output_image
    )

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