import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import os
import logging
import cv2
import easyocr
from plate_utils import extract_plate, standardize_plate

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

helmet_model = YOLO("Helmet-Detect-model/best.pt")
reader = easyocr.Reader(['en'])

def detect_helmet(image):
    results = helmet_model(image)
    results_plotted = results[0].plot()
    return Image.fromarray(results_plotted)

def process_plate(image):
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    plate_img = extract_plate(image)
    if plate_img is None:
        return "No plate detected", None
    
    results = reader.readtext(plate_img)
    text = ""
    for (bbox, text_result, prob) in results:
        text += text_result + " "
    
    text = text.strip()
    standardized = standardize_plate(text)
    if standardized:
        text = f"Raw: {text}\nStandardized: {standardized}"
    
    return text, Image.fromarray(plate_img)

def get_example_images(folder_path):
    examples = []
    if os.path.exists(folder_path):
        for filename in os.listdir(folder_path):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')) and not filename.startswith('.'):
                examples.append(os.path.join(folder_path, filename))
    return examples

with gr.Blocks(
    title="Helmet Detection and Number Plate OCR",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")
) as demo:
    gr.Markdown(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <h1 style="color: #2196F3; font-size: 2.5rem; margin-bottom: 1rem;">
                Helmet Detection & Number Plate Recognition
            </h1>
            <p style="font-size: 1.1rem; color: #666; margin-bottom: 2rem;">
                Created by Abhay Gupta | Advanced Computer Vision Project
            </p>
        </div>
        """
    )
    
    with gr.Tabs() as tabs:
        with gr.Tab("Helmet Detection", id=1):
            with gr.Row():
                with gr.Column(scale=1):
                    helmet_input = gr.Image(
                        type="pil",
                        label="Upload Image",
                        tool="select",
                        height=300
                    )
                    helmet_button = gr.Button(
                        "Detect Helmets",
                        variant="primary",
                        scale=1
                    )
                
                with gr.Column(scale=1):
                    helmet_output = gr.Image(
                        type="pil",
                        label="Detection Result",
                        height=300
                    )
            
            gr.Markdown("### Example Images")
            helmet_examples = get_example_images("Helmet-Detect-model/test_images")
            if helmet_examples:
                gr.Examples(
                    examples=helmet_examples,
                    inputs=helmet_input,
                    outputs=helmet_output,
                    fn=detect_helmet,
                    cache_examples=True,
                    label="Click to try an example"
                )
        
        with gr.Tab("Number Plate OCR", id=2):
            with gr.Row():
                with gr.Column(scale=1):
                    plate_input = gr.Image(
                        type="pil",
                        label="Upload Image",
                        tool="select",
                        height=300
                    )
                    plate_button = gr.Button(
                        "Read Number Plate",
                        variant="primary",
                        scale=1
                    )
                
                with gr.Column(scale=1):
                    with gr.Group():
                        plate_text = gr.Textbox(
                            label="Extracted Text",
                            lines=3,
                            show_copy_button=True
                        )
                        plate_output = gr.Image(
                            type="pil",
                            label="Processed Plate",
                            height=200
                        )
            
            gr.Markdown("### Example Images")
            plate_examples = get_example_images("Number_plate_OCR/test_images")
            if plate_examples:
                gr.Examples(
                    examples=plate_examples,
                    inputs=plate_input,
                    outputs=[plate_text, plate_output],
                    fn=process_plate,
                    cache_examples=True,
                    label="Click to try an example"
                )
    
    gr.Markdown(
        """
        <div style="text-align: center; margin-top: 2rem; padding: 1rem; background: #f5f5f5; border-radius: 8px;">
            <p style="color: #666; margin: 0;">
                Computer Vision Project by Abhay Gupta
            </p>
        </div>
        """
    )
    
    helmet_button.click(
        fn=detect_helmet,
        inputs=helmet_input,
        outputs=helmet_output
    )
    
    plate_button.click(
        fn=process_plate,
        inputs=plate_input,
        outputs=[plate_text, plate_output]
    )

demo.launch()