Update app.py

app.py

@@ -1,36 +1,69 @@
-import gradio as gr
-import os
 import torch
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import cv2
+import re
 from PIL import Image
+import gradio as gr
+import numpy as np
+
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='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.")
 
-                           
-#subprocess.run(["mv","content/custom_data.yaml","./yolov5/data"])         
-
-
-def load_model():
- '''
- Loading hub model & setting the preferences for the model  
- '''
- model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolo-v5.pt')
- model.conf = 0.38 
- model.dnn=True
- model.agnostic=True
- return model
-
-model=load_model()
-#, force_reload=True
-def detect(inp):
- #g = (size / max(inp.size))  #gain
- #im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize 
- results = model(inp,size=640)  # inference
- results.render()  # updates results.imgs with boxes and labels
- return Image.fromarray(results.ims[0])
- 
-
-inp = gr.Image(type="pil", label="Original Image")
-output = gr.Image(type="pil", label="Output Image")
-
-
-io=gr.Interface(fn=detect, inputs=inp, outputs=output, title='Party Symbol Detection',examples=['Content/4.jpg','Content/10.jpg','Content/18.jpg'],theme='peach')
-io.launch(debug=True,share=False)
-  
+interface.launch()