import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import cv2
import re
from PIL import Image
import gradio as gr
import numpy as np
import yolov5

model = yolov5.load('yolo-v5.pt')


model.conf = 0.80 

processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-printed')
ocr = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-printed')

def extract_coordinates(img, model):
    results = model(img)
    cordinates = results.xyxy[0][:, :-1]
    return cordinates

def read_plate_number(results, frame, cordinates):
    plate_numbers = []
    n = len(results)

    for i in range(n):
        row = cordinates[i]
        if row[4] >= 0.5:
            xmin, ymin, xmax, ymax = map(int, row[:4])
            plate = frame[ymin:ymax, xmin:xmax]

            pixel_values = processor(images=plate, return_tensors="pt").pixel_values
            generated_ids = ocr.generate(pixel_values)
            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

            cleaned_text = clean_plate_number(generated_text)
            plate_numbers.append(cleaned_text)
            
    return plate_numbers

def clean_plate_number(text):
    cleaned_text = re.sub(r'[^a-zA-Z0-9]', '', text)
    
    if any(char.isalpha() for char in cleaned_text) and any(char.isdigit() for char in cleaned_text):
        plate_number = cleaned_text[-7:]
        return plate_number
    
    return ""

def perform_ocr_on_image(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    results = model(img)
    cordinates = extract_coordinates(img, model)
    
    if len(cordinates) == 0:
        return "Nenhuma placa encontrada."
    
    plate_number = read_plate_number(results.pred[0], img, cordinates)
    
    if plate_number:
        return plate_number[0].lower()
    else:
        return "Não foi possível reconhecer a placa."

interface = gr.Interface(fn=perform_ocr_on_image, 
                         inputs=gr.Image(type="pil"),
                         outputs="text",
                         title="Reconhecimento de Placas de Automóveis",
                         description="Envie uma imagem e receba o número da placa.")

interface.launch()