Delete vehicle_attribute_v5.py

vehicle_attribute_v5.py

@@ -1,128 +0,0 @@
-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:
-        # Ensure all rectangle coordinates are integers
-        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])
-        # Calculate the coordinates for the rectangle
-        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)
-        # Draw the rectangle and text
-        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)
-        print("Text: ", text)
-        print("Coordinates: ", top_left, bottom_right)
-        return img0  # Return the image unchanged in case of an error
-
-# 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
-
-if __name__ == "__main__":
-    image_path = r"./canada"  #
-    output_path = './result_new_img'  #
-    file_list = []
-    allFilePath(image_path, file_list)
-    if not os.path.exists(output_path):
-        os.mkdir(output_path)
-    save_path = output_path
-    count = 0
-    sumTime = 0
-
-    pedestrain_predictor = Predict()
-
-    for pic_ in file_list:
-        if not (pic_.endswith(".jpg") or pic_.endswith(".png")):
-            continue
-        count += 1
-        time_begin = time.time()
-        img = cv_imread(pic_)
-        if img.shape[-1] == 4:
-            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
-        img0 = copy.deepcopy(img)
-        result_list = pedestrain_predictor.predict(img)
-        time_end = time.time()
-        time_gap = (time_end - time_begin)
-        sumTime += time_gap
-        print(count, pic_, time_gap)
-        img0 = draw_result(img0, result_list)
-        img_name = os.path.basename(pic_)
-        save_img_path = os.path.join(save_path, img_name)
-        cv2.imencode('.jpg', img0)[1].tofile(save_img_path)
-
-    print(f"cost time  {sumTime}, {sumTime / (len(file_list)) * 1000} ms per img")