from typing import get_args
from PIL import Image
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile
import os
import subprocess

from fast_alpr import ALPR
from fast_alpr.default_detector import PlateDetectorModel
from fast_alpr.default_ocr import OcrModel

# Loading YOLO model
model = YOLO('best.pt') 
img_dim = (640, 640)

# Default models for plate recognition
DETECTOR_MODELS = list(get_args(PlateDetectorModel))
OCR_MODELS = list(get_args(OcrModel))
# Put global OCR first
OCR_MODELS.remove("global-plates-mobile-vit-v2-model")
OCR_MODELS.insert(0, "global-plates-mobile-vit-v2-model")

def predict(image, conf_threshold, iou_threshold):
    # Resizing
    image = image.resize(img_dim)
    
    # Convert from PIL to OpenCV format
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Run inference
    results = model.predict(
        image, 
        conf=conf_threshold, 
        iou=iou_threshold,
        imgsz=640
    )
    
    # Get annotated image
    annotated_image = results[0].plot()
    
    # Perform OCR on detected objects
    ocr_text= ""
    for box in results[0].boxes.xyxy.cpu().numpy():
        x1, y1, x2, y2 = map(int, box)
        cropped = image[y1:y2, x1:x2]
        
        # Skip if the cropped region is too small
        if cropped.size == 0:
            continue
        # Apply detector for plate region 
        alpr = ALPR(detector_model=DETECTOR_MODELS[0], ocr_model=OCR_MODELS[0])
        alpr_results = alpr.predict(cropped)

        if alpr_results:
            res = alpr_results[0]
            # Access the detection and OCR attributes from ALPRResult
            plate_text = res.ocr.text if res.ocr else "N/A"
            plate_confidence = res.ocr.confidence if res.ocr else 0.0
            ocr_text += f"- Detected Plate: {plate_text} with confidence {plate_confidence:.2f}\n"
        
            # Add text to annotated image
            cv2.putText(annotated_image, plate_text, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
    # Convert back to RGB for Gradio display
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)     
    
    return annotated_image, ocr_text
    

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# MIA-Yolov8 for peruvian plate recognition")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil")
            conf_slider = gr.Slider(0, 1, value=0.25, label="Confidence Threshold")
            iou_slider = gr.Slider(0, 1, value=0.45, label="IOU Threshold")
            submit_btn = gr.Button("Run model")
        with gr.Column():
            output_image = gr.Image(label="Detected Objects")
            ocr_output = gr.Textbox(label="OCR Results")
    
    submit_btn.click(
        fn=predict,
        inputs=[input_image, conf_slider, iou_slider],
        outputs=[output_image, ocr_output]
    )

demo.launch()