Upload 3 files

app.py

@@ -0,0 +1,123 @@
+import gradio as gr
+from PIL import Image as PILImage 
+import onnxruntime
+import numpy as np
+import cv2
+import copy
+import os
+import time
+from fonts.cv_puttext import cv2ImgAddText
+
+from car_plate.yolov5_plate_onnx_infer import init_car_plate_detect_model, init_car_plate_rec_model, detect_plate, rec_plate
+import re
+
+import json
+
+def cv_imread(path):
+    img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
+    return img
+
+def allFilePath(rootPath, allFileList):
+    fileList = os.listdir(rootPath)
+    for temp in fileList:
+        if os.path.isfile(os.path.join(rootPath, temp)):
+            allFileList.append(os.path.join(rootPath, temp))
+        else:
+            allFilePath(os.path.join(rootPath, temp), allFileList)
+
+def draw_car_attribute(text, img0, n, rect):
+    try:
+        rect = [int(x) for x in rect]
+        labelSize = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        sp_size = labelSize[1] + round(1.6 * labelSize[0][1])
+        top_left = (rect[0], int(rect[1] + n * sp_size))
+        bottom_right = (int(rect[0] + round(1.4 * labelSize[0][0])), rect[1] + (n + 1) * sp_size)
+        img0 = cv2.rectangle(img0, top_left, bottom_right, (255, 255, 255), cv2.FILLED)
+        img0 = cv2ImgAddText(img0, text, rect[0], int(rect[1] + n * sp_size), (0, 0, 0), 18)
+        return img0
+    except Exception as e:
+        print("Error drawing car attribute: ", e)
+        return img0
+
+# Ensure this modified function is used in the draw_result function.
+
+def draw_result(img0, result_list):
+    for result in result_list:
+        if not result:
+            continue
+        n = 0
+        rect = result.get('rect')
+        if rect and len(rect) == 4:  # Ensure rect has four elements
+            cv2.rectangle(img0, (int(rect[0]), int(rect[1])), (int(rect[2]), int(rect[3])), (255, 0, 0), 2)
+            if 'plate' in result and len(result['plate']) > 0:
+                for plate_ in result['plate']:
+                    plate_no = plate_.get('plate_no')
+                    if plate_no and len(plate_no) > 0:
+                        result_plate = "PLATE NUMBER: " + plate_no
+                        plate_rect = plate_.get('rect')
+                        if plate_rect and len(plate_rect) == 4:  # Check if plate_rect is correctly formed
+                            cv2.rectangle(img0, (int(plate_rect[0]), int(plate_rect[1])), (int(plate_rect[2]), int(plate_rect[3])), (0, 255, 0), 2)
+                            img0 = draw_car_attribute(result_plate, img0, n, plate_rect)
+                            n += 1
+        else:
+            print("Invalid 'rect' data:", rect)
+    return img0
+
+class Predict:
+    def __init__(self):
+        providers = ['CPUExecutionProvider']
+        car_plate_detect_model_path = r"weights/best_exp3.onnx"  # plate detect onnx
+        car_plate_rec_model_path = r"weights/plate_rec_color_0820.onnx"  # plate recognition onnx
+        self.session_detect_plate = init_car_plate_detect_model(car_plate_detect_model_path, providers)  #
+        self.session_rec_plate = init_car_plate_rec_model(car_plate_rec_model_path, providers)  #
+
+    def predict(self, img_data):
+        img0 = copy.deepcopy(img_data)
+        result_list = []
+        outputs = detect_plate(img0, self.session_detect_plate, 640)
+        plate_list = rec_plate(outputs, img0, self.session_rec_plate)
+        plate_dict_list = []
+        for result in plate_list:
+            plate_dict = {}
+            plate_no = result['plate_no']
+            plate_rect = result['rect']
+            plate_lands = result['landmarks']
+            plate_dict['plate_no'] = plate_no
+            plate_dict['rect'] = plate_rect
+            plate_dict['landmarks'] = plate_lands
+            plate_dict_list.append(plate_dict)
+        if plate_dict_list:
+            result_list.append({'rect': plate_rect, 'plate': plate_dict_list})
+        return result_list
+
+# 在使用 car_plate_demo 之前实例化 Predict 类
+car_plate_predictor = Predict()
+
+# Define a function to wrap the prediction and visualization
+def car_plate_demo(input_image):
+    # Convert the PIL image to an OpenCV image
+    input_image = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    
+    # Call the predictor to get the results
+    result_list = car_plate_predictor.predict(input_image)
+    
+    # Draw the results on the image
+    output_image = draw_result(copy.deepcopy(input_image), result_list)
+    
+    # Convert the OpenCV image back to PIL to display in Gradio
+    output_image_pil = PILImage.fromarray(cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB))
+    
+    return output_image_pil
+
+demo = gr.Interface(
+    fn=car_plate_demo,
+    inputs=gr.Image(label="Upload your car image"),
+    outputs=gr.Image(label="Processed Image"),
+    examples=["whitecar.jpg", "low-light-black-car.jpg"]
+)
+
+
+
+# Run the Gradio app
+if __name__ == "__main__":
+    demo.launch()

packages.txt

@@ -0,0 +1 @@
+python3-opencv

requirements.txt

@@ -0,0 +1,10 @@
+gradio
+numpy
+opencv-python-headless
+Pillow
+onnxruntime
+opencv-python
+jinja2
+numpy
+matplotlib
+scikit-learn