import numpy as np
import gradio as gr
from PIL import Image

from fast_alpr import ALPR

# Fixed model
DETECTOR_MODEL = "yolo-v9-s-608-license-plate-end2end"

def alpr_inference(image):
    if image is None:
        return None, "Please upload an image to continue."
    img = image.convert("RGB")
    img_array = np.array(img)
    alpr = ALPR(detector_model=DETECTOR_MODEL)
    # Only detect plates, do not perform OCR
    results = alpr.predict(img_array)  # Use predict() as detect() is not available
    
    # Draw only bounding boxes, not text
    annotated_img_array = img_array.copy()
    annotated_img = Image.fromarray(annotated_img_array)
    from PIL import ImageDraw
    draw = ImageDraw.Draw(annotated_img)
    for result in results:
        detection = getattr(result, 'detection', None)
        if detection is not None:
            bbox_obj = getattr(detection, 'bounding_box', None)
            if bbox_obj is not None:
                bbox = [int(bbox_obj.x1), int(bbox_obj.y1), int(bbox_obj.x2), int(bbox_obj.y2)]
                draw.rectangle(bbox, outline="red", width=3)
    if results:
        detection_results = f"Detected {len(results)} license plate(s)."
    else:
        detection_results = "No license plate detected 😔."
    return annotated_img, detection_results

with gr.Blocks() as demo:
    gr.Markdown("# License Plate Detection")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload an image of a vehicle with a license plate")
            submit_btn = gr.Button("Run Plate Detection")
        with gr.Column():
            annotated_output = gr.Image(label="Annotated Image with Plate Detection")
            detection_output = gr.Markdown(label="Detection Results")
    submit_btn.click(
        alpr_inference,
        inputs=[image_input],
        outputs=[annotated_output, detection_output]
    )

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