Spaces:
Running
Running
| 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", | |
| examples=['1.jpg','2.jpg'], | |
| description="Envie uma imagem e receba o número da placa.") | |
| interface.launch() |