Update app.py

app.py

@@ -1,69 +1,36 @@
+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.")
 
-interface.launch()
+                           
+#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)
+