import gradio as gr
from solution import TrafficViolationDetector

# Initialize globally but lazily to avoid Hugging Face 60s boot timeout
detector = None

def detect_violations(image_path):
    global detector
    if image_path is None:
        return {"error": "No image provided"}
    
    if detector is None:
        print("Lazy loading models on first request...")
        detector = TrafficViolationDetector(model_dir="./models")
        print("Models loaded successfully!")
        
    try:
        # The detector.predict expects a path to the image
        result = detector.predict(image_path)
        return result
    except Exception as e:
        return {"error": str(e)}

# Create the Gradio interface
iface = gr.Interface(
    fn=detect_violations,
    inputs=gr.Image(type="filepath", label="Upload Traffic Image"),
    outputs=gr.JSON(label="Violation Results"),
    title="Traffic Rule Violation Detection API",
    description="Upload an image to detect traffic violations. Supports two-wheelers (helmet, over-riding, wrong-way) and four-wheelers (seatbelt, wrong-way). Detects and runs OCR on the license plates of violating vehicles.\n\nThis application can be accessed programmatically via its built-in API.",
)

if __name__ == "__main__":
    # Launch on 0.0.0.0 to allow Hugging Face to route traffic
    iface.launch(server_name="0.0.0.0", server_port=7860)