增加motorcycle，bus和truck检测，优化车牌检测，增加单元测试

README.md

@@ -5,4 +5,33 @@
 - 如果目标类别是car，则用ocr提取车牌号；
 - 如果目标类别是person，先用交警模型判断是否是交警，如果不是交警，则对人脸图像清晰度大于等于一定阈值的，记录该人脸的embedding；
 
-返回每一帧处理得到的所有车牌号列表（去重）和所有人脸embedding列表；
+返回每一帧处理得到的所有车牌号列表（去重）和所有人脸embedding列表；
+
+## 性能测试
+| 视频名 |   显存   | gpu使用率 | 处理帧数 |  耗时   | half  | 跳帧 | 检测结果（人脸/车牌数） |
+| :----: | :------: | :------: | :------: | :-----: | :---: | :--: | :---------------------: |
+| 2.MP4  | 2.0G左右 |   50%+    |  26975   | 166.57s | False |  -   |          145/0          |
+| 2.MP4  | 1.9G左右 |   50%+    |  26975   | 162.65s | True  |  -   |          145/0          |
+| 2.MP4  | 2.0G左右 |   40%+    |   2997   | 28.79s  | False |  9   |          17/0           |
+| 2.MP4  | 1.9G左右 |   40%+    |   2997   | 28.83s  | True  |  9   |          17/0           |
+| 1.MP4  | 2.0G左右 |   50%+    |  26991   | 178.69s | False |  -   |          578/0          |
+| 1.MP4  | 2.0G左右 |   50%+    |  26991   | 176.29s | True  |  -   |          578/0          |
+| 1.MP4  | 2.0G左右 |   40%+    |   2999   | 30.12s  | False |  9  |          60/0           |
+| 1.MP4  | 2.0G左右 |   40%+    |   2999   | 30.04s  | True  |  9  |          60/0           |
+
+### float16推理
+- 如果不跳帧，半精度推理后速度略有提升；
+- 如果本身处理帧数少，速度没啥影响；
+- 对检测结果也几乎没有影响；
+
+### 跳帧
+- 相比不跳帧，耗时明显减少，筛选得到的结果也明显减少；
+- gpu使用率降低；
+
+### 更新
+- 20240125：对被执法车加入truck，motorcycle，bus类别，增加摩托车车牌识别规则（车牌内容格式为两行）；
+
+### 部分视频无法检测人或车原因
+- 人脸相似个数太少，导致合并时间段时被过滤；
+- yolo对车目标检测置信度太低，或直接没有检测到，导致不会进行下一步ocr检测；
+- 车牌不够清晰导致ocr检测无结果；

src/algs/ocr_alg_test.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Filter_Object 
+@File ：ocr_alg_test.py
+@IDE ：PyCharm 
+@Author ：zjj
+@Date ：2024/1/23 13:26 
+'''
+import unittest
+
+import cv2
+
+from algs.ocr_alg import ocr_predict
+
+
+class OcrTest(unittest.TestCase):
+
+    @staticmethod
+    def test_ocr_predict():
+        frame = cv2.imread('../data/truck3.png')
+        result, vis_im = ocr_predict(frame)
+        print(result)
+        cv2.imwrite('../data/ocr_result.png', vis_im)

src/algs/yolo_alg.py

@@ -1,59 +0,0 @@
-import fastdeploy as fd
-import cv2
-import os
-
-
-def parse_arguments():
-    import argparse
-    import ast
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", default=None, help="Path of yolov8 model.")
-    parser.add_argument(
-        "--image", default=None, help="Path of test image file.")
-    parser.add_argument(
-        "--device",
-        type=str,
-        default='gpu',
-        help="Type of inference device, support 'cpu' or 'gpu' or 'kunlunxin'.")
-    parser.add_argument(
-        "--use_trt",
-        type=ast.literal_eval,
-        default=True,
-        help="Wether to use tensorrt.")
-    return parser.parse_args()
-
-
-def build_option(args):
-    option = fd.RuntimeOption()
-
-    if args.device.lower() == "gpu":
-        option.use_gpu()
-
-    if args.device.lower() == "ascend":
-        option.use_ascend()
-
-    if args.use_trt:
-        option.use_trt_backend()
-        option.set_trt_input_shape("images", [1, 3, 640, 640])
-    return option
-
-
-args = parse_arguments()
-
-# Configure runtime, load model
-runtime_option = build_option(args)
-model = fd.vision.detection.YOLOv8(args.model, runtime_option=runtime_option)
-
-def yolo_predict(im):
-# # Predicting image
-# if args.image is None:
-#     image = fd.utils.get_detection_test_image()
-# else:
-#     image = args.image
-# im = cv2.imread(image)
-    result = model.predict(im)
-    # Visualization
-    vis_im = fd.vision.vis_detection(im, result)
-    # cv2.imwrite("visualized_result.jpg", vis_im)
-    # print("Visualized result save in ./visualized_result.jpg")
-    return result,vis_im

src/algs/yolo_alg_bk.py

@@ -1,41 +0,0 @@
-import fastdeploy as fd
-from pathlib import Path
-import cv2
-import os
-
-
-class YOLOAlg:
-    def __init__(self, model_path) -> None:
-        super(YOLOAlg, self).__init__()
-        self.model_path = model_path
-        self.model = self.init_model()
-
-    def build_option(self):
-        option = fd.RuntimeOption()
-        option.use_gpu()
-
-        option.use_trt_backend()
-        option.set_trt_input_shape("images", [1, 3, 640, 640])
-
-        trt_path = Path(self.model_path).with_suffix(".trt")
-        option.set_trt_cache_file(trt_path.as_posix())
-
-        return option
-    
-    def init_model(self):
-        # Configure runtime, load model
-        runtime_option = self.build_option()
-        model = fd.vision.detection.YOLOv8(self.model_path, runtime_option=runtime_option)
-        return model
-
-    def predict_yolo(self, bgr_img):
-        result = self.model.predict(bgr_img)
-
-        rendered_img = bgr_img.copy()
-
-        # Visualization
-        vis_im = fd.vision.vis_detection(rendered_img, result)
-        # cv2.imwrite("visualized_result.jpg", vis_im)
-        # print("Visualized result save in ./visualized_result.jpg")
-
-        return result, vis_im

src/config.py

@@ -13,7 +13,12 @@ from pathlib import Path
 
 PERSON = "person"
 CAR = "car"
-LABEL_NAMES = {0: PERSON, 2: CAR}
+TRUCK = 'truck'
+MOTORCYCLE = 'motorcycle'
+BUS = 'bus'
+MOTOR_VEHICLE = 'car_related'
+MOTOR_VEHICLE_LIST = [CAR, TRUCK, MOTORCYCLE, BUS]
+LABEL_NAMES = {0: PERSON, 2: CAR, 3: MOTORCYCLE, 5: BUS, 7: TRUCK}
 
 # 过程中保存图片路径
 # SAVE_BASE_DIR = Path('/home/TP/PMP/media/tp_result')
@@ -69,11 +74,11 @@ TOPN = {PERSON: 5, CAR: 2}
 
 
 # 人脸检测阈值 0.85
-FACE_THRESHOLD = 0.85
+FACE_THRESHOLD = 0.83
 # 目标检测置信度阈值 0.8
-DETECT_CONF_THRESHOLD = 0.8
+DETECT_CONF_THRESHOLD = 0.6
 # 框面积阈值 todo 存在有的被执法人就是离镜头远或者模型画的框小 15
-BOX_AREA_THRESHOLD = 15
+BOX_AREA_THRESHOLD = 5
 
 # 每一帧内处理面积top2个大框图
 TOPN_AREA = 2
@@ -83,3 +88,6 @@ POLICE_IOU = 0.5
 
 # 人脸清晰度阈值 0-100之间
 ARTICULATION_THD = 30
+
+# 摩托车两行车牌文字box的接近程度（上面box的ymax和下面box的ymin差值绝对值）
+DELTA_Y = 10

src/fd_face_detection.py

@@ -25,6 +25,7 @@ class FaceRecognition:
         det_model=SCRFD_PATH_DICT["onnx"],
         device="gpu",
         use_trt=False,
+        half=False,
     ):
         """
         初始化相机人脸检测器
@@ -32,11 +33,13 @@ class FaceRecognition:
         :param det_model: 人脸检测模型
         :param device: 使用设备"cpu"/"gpu"
         :param use_trt: 是否使用tenserRT加速
+        :param half: 是否半精度推理
         """
         self.id_model = id_model
         self.det_model = det_model
         self.device = device
         self.use_trt = use_trt
+        self.half = half
 
         # 用统一的方法加载模型
         self.det_model = self.load_model(fd.vision.facedet.SCRFD, det_model, "SCRFD")
@@ -98,6 +101,9 @@ class FaceRecognition:
             option.use_gpu()
         if self.use_trt:
             option.use_trt_backend()
+            if self.half:
+                option.trt_option.enable_fp16 = True
+
             if model_type == "PartialFC":
                 option.set_trt_input_shape("data", [1, 3, 112, 112])  # fd默认参数
                 option.set_trt_cache_file(INSIGHTFACE_PATH_DICT["trt"])
@@ -194,20 +200,16 @@ class FaceRecognition:
 if __name__ == "__main__":
     # faceid = FaceRecognition(device="cpu", use_trt=False)
     faceid = FaceRecognition(
-        id_model=r"E:\resources\DT3 models\partial_fc_glint360k_r100.onnx",
-        det_model=r"E:\resources\DT3 models\scrfd_500m_bnkps_shape640x640.onnx",
+        id_model=r"/home/user/zhoujinjuan/models/partial_fc_glint360k_r100.onnx",
+        det_model=r"/home/user/zhoujinjuan/models/scrfd_500m_bnkps_shape640x640.onnx",
         device="cpu",
         use_trt=False,
     )
-    # img = cv2.imread('../data/track/13.png')
-    # embed = faceid.get_img_embedding(img)
-    # has_face, face = faceid.get_crop_face(img)
-    # if has_face:
-    #     cv2.imwrite('../data/track/4_face.png', face)
-    # else:
-    #     print('no face')
-
-    # img1 = cv2.imread(r'E:\\tmp\198_person\198_0.png')
-    # img2 = cv2.imread(r'E:\\tmp\198_person\198_1.png')
-    # cos = faceid.face_compare(img2, img1)
-    # print(cos)
+    img = cv2.imread('data/11.png')
+    embed = faceid.get_img_embedding(img)
+    has_face, face = faceid.get_crop_face(img)
+    if has_face:
+        cv2.imwrite('data/face.png', face)
+    else:
+        print('no face')
+

src/fd_yolo.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Filter_Object 
+@File ：fd_yolo.py
+@IDE ：PyCharm 
+@Author ：zjj
+@Date ：2024/1/19 10:45 
+'''
+import fastdeploy as fd
+from loguru import logger
+
+from config import POLICE_PATH
+
+
+class FdYolov8(object):
+    def __init__(self, model_path=POLICE_PATH, device="gpu", use_trt=False, half=False):
+        self.model_path = model_path
+        self.device = device
+        self.use_trt = use_trt
+        self.half = half
+        self.model = None
+        # 加载模型
+        self.init_model()
+
+    def build_option(self):
+        """
+        构建运行时选项
+        :return: option
+        """
+        option = fd.RuntimeOption()
+        if self.device.lower() == "gpu":
+            option.use_gpu()
+        if self.use_trt:
+            option.use_trt_backend()
+            if self.half:
+                option.trt_option.enable_fp16 = True
+        return option
+
+    def init_model(self):
+        runtime_option = self.build_option()
+        self.model = fd.vision.detection.YOLOv8(
+            self.model_path, runtime_option=runtime_option
+        )
+        logger.info('fastdeploy yolo模型加载成功！')
+
+    def predict(self, bgr_img):
+        result = self.model.predict(bgr_img)
+        return result

src/main_v2.py

@@ -7,10 +7,6 @@
 @Author ：zjj
 @Date ：2023/9/26 9:29 
 """
-import uuid
-
-import cv2
-import fastdeploy as fd
 import numpy as np
 from loguru import logger
 from ultralytics import YOLO
@@ -23,9 +19,10 @@ from config import (
     LABEL_NAMES,
     BOX_AREA_THRESHOLD,
     DETECT_CONF_THRESHOLD,
-    POLICE_PATH, ARTICULATION_THD,
+    ARTICULATION_THD, MOTOR_VEHICLE, MOTOR_VEHICLE_LIST,
 )
 from fd_face_detection import FaceRecognition
+from fd_yolo import FdYolov8
 from utils import (
     extract_yolo_results,
     get_police,
@@ -38,47 +35,47 @@ from utils import (
 
 
 class TrackMain(object):
-    def __init__(self, detect_model_path=SEGMENT_MODEL, device="gpu", use_trt=False):
-        # fastdeploy调用GPU时会使用
-        self.option = fd.RuntimeOption().use_gpu()
+    def __init__(self, detect_model_path=SEGMENT_MODEL, device="gpu", use_trt=False, half=False):
         # 目标检测模型
-        self.model = YOLO(detect_model_path)
+        self.yolo_seg_model = YOLO(detect_model_path)
         # 人脸检测和识别对象
-        self.face_det = FaceRecognition(device=device, use_trt=use_trt)
+        self.face_det = FaceRecognition(device=device, use_trt=use_trt, half=half)
         # 交警模型
-        # self.model_traffic_police = YOLO(POLICE_PATH)
-        self.model_traffic_police = fd.vision.detection.YOLOv8(
-            POLICE_PATH, runtime_option=self.option # 使用fastdeploy加载交警模型
-        )
+        self.model_traffic_police = FdYolov8(device=device, use_trt=use_trt, half=half)
 
     def process_one_frame(self, frame):
-        results = self.model(frame)
+        """
+        处理一帧数据
+        :param frame:
+        :return:
+        """
+        results = self.yolo_seg_model(frame)
+
         # 提取模型检测结果 infos=[xywhs, cls, scores, xy_list]
         is_hit, infos = extract_yolo_results(results)
         if not is_hit:
             return [], []
 
-        # before:没有筛选之前统计目标频率，保存截图时筛选top2大框
-        # now:不统计频率，大框筛选之后，是否有必要取topN再过滤？？---先设置过滤，不想过滤则可以调整N值很大
-        # I：面积和置信度过滤，统计剩余目标的相关信息：左上角和右下角坐标,label,conf,box_area,轮廓坐标
+        # 不统计频率，大框筛选之后，设置TOPN过滤，不想过滤则可以调整N值很大
+        # 面积和置信度过滤，统计剩余目标的相关信息：左上角和右下角坐标,label,conf,box_area,轮廓坐标
         tinfos_each_frame = self.statistics_one_frame(infos)
 
-        # 每帧每个类别显示topn大的目标，按文件夹保存，person类别保留人脸，相似的保留一个，car类别保存目标检测框图，有一个算一个；
         # annotated_frame = results[0].plot()
-        face_embeddings, licenses = self.process_topn_in_one_frame(frame, tinfos_each_frame)
+        face_embeddings, licenses, _ = self.process_topn_in_one_frame(frame, tinfos_each_frame)
         return face_embeddings, licenses
 
-    def statistics_one_frame(self, infos):
+    @classmethod
+    def statistics_one_frame(cls, infos):
         """
         每一帧内统计每个目标的信息
         :param infos:
         :return:
         """
-        xywhs, cls, scores, xy_list = infos
-        tinfos_each_frame = {PERSON: [], CAR: []}
+        xywhs, labels, scores, xy_list = infos
+        tinfos_each_frame = {PERSON: [], MOTOR_VEHICLE: []}
 
         for xywh, label_id, score, polygon in zip(
-            xywhs, cls, scores, xy_list
+            xywhs, labels, scores, xy_list
         ):
             if label_id not in list(LABEL_NAMES.keys()):
                 continue
@@ -96,6 +93,7 @@ class TrackMain(object):
             conf = round(score, 2)
 
             # 通过面积过滤掉一些小框
+            # todo 预筛选，阈值不能太高，后面还会有topN过滤
             if s <= BOX_AREA_THRESHOLD:
                 continue
             if conf <= DETECT_CONF_THRESHOLD:
@@ -107,13 +105,19 @@ class TrackMain(object):
             info['conf'] = conf
             info['box_area'] = s
             info['polygon_indexs'] = polygon.astype(int)
-            tinfos_each_frame[label].append(info)
+            # 机动车包含car,truck,motorcycle,bus
+            if label in MOTOR_VEHICLE_LIST:
+                tinfos_each_frame[MOTOR_VEHICLE].append(info)
+            else:
+                tinfos_each_frame[label].append(info)
 
         return tinfos_each_frame
 
     def process_topn_in_one_frame(self, frame, tinfos_each_frame):
         """
-        处理每一帧中的topn大框对象
+        处理每一帧中的topn大框对象:
+        person类：先用交警模型过滤掉交警，获取目标的人脸，根据清晰度过滤，对保留下来的人脸提取embedding；
+        car类：ocr识别车牌号，每一帧中车牌号去重；
         :param frame:
         :param tinfos_each_frame: 每个类别下各个目标的相关信息
         :return:
@@ -121,14 +125,16 @@ class TrackMain(object):
         frame_copy = frame.copy()
 
         police_indexs = []
+        # 存储一帧中所有车牌号
         all_licenses = []
+        # 存储一帧中所有人脸embedding
         all_face_embeddings = []
-        for label, target_infos in tinfos_each_frame.items():
+        for label_alias, target_infos in tinfos_each_frame.items():
             # 按照框面积大小降序排序
             target_infos.sort(key=lambda x: x['box_area'], reverse=True)
 
             # todo 交警 先对整个图片上前topn的大框图中的person，送到交警检测模型中检测是否是交警
-            if label == PERSON:
+            if label_alias == PERSON:
                 police_indexs = get_police(
                     frame_copy, target_infos[:TOPN_AREA], self.model_traffic_police
                 )
@@ -144,8 +150,7 @@ class TrackMain(object):
                 target_img = frame[p1[1]: p2[1], p1[0]: p2[0]]
                 target_img = target_img.astype(np.uint8)
 
-                if label == CAR:
-                    # todo 调用ocr处理车牌号，统计车牌出现次数，保存有车牌的图片
+                if label_alias == MOTOR_VEHICLE:
                     # licenses = predict_ocr(target_img, self.ocr)
                     licenses = predict_ocr(target_img)
                     licenses = list(set(licenses))
@@ -153,14 +158,13 @@ class TrackMain(object):
                     if licenses:
                         is_hit = True
 
-                elif label == PERSON:
+                elif label_alias == PERSON:
                     # 是交警，则不处理
                     if index in police_indexs:
                         continue
-                    # 提取目标的语义分割图像，配白底，保存； 保存所有，不考虑相似的只保存一个，防止有车牌的被过滤掉
+                    # 二值图像，人-1， 背景-0
                     target_mask = get_target_mask(frame, polygon)
-                    # person类型，检测并提取人脸，相似的只保存一个，即保存的时候就过滤一些
-                    # todo 对整个图检测人脸可能得到多个人脸，但只对target_img的白底目标图检测人脸，很可能无法检测到人脸，因此结合目标检测的框和人脸检测算法，框定目标的人脸
+                    # todo 对整个图检测人脸可能得到多个人脸，但只对target_img的白底目标图检测人脸，很可能无法检测到人脸，因此结合实例mask定位目标人脸
                     has_face, face = self.face_det.get_target_face(frame, target_mask)
                     if has_face:
                         # 用人脸清晰度过滤
@@ -172,40 +176,10 @@ class TrackMain(object):
                         face_embedding = self.face_det.get_img_embedding(face)
                         all_face_embeddings.append(face_embedding)
                         is_hit = True
-                # todo test
+                # test
                 if is_hit:
                     frame_copy = draw_rectangle_text(
-                        frame_copy, index, p1, p2, label, info['conf'], -1, info['box_area']
+                        frame_copy, index, p1, p2, label_alias, info['conf'], -1, info['box_area']
                     )
-
-        # todo test 测试过程中保存图片可视化
-        # if all_face_embeddings and not all_licenses:
-        #     name = str(uuid.uuid4())[:7]
-        #     cv2.imwrite('/mnt/large/zhoujinjuan_data/data/result-3/face_' + name + '.png', frame_copy)
-        # elif not all_face_embeddings and all_licenses:
-        #     name = str(uuid.uuid4())[:7]
-        #     cv2.imwrite('/mnt/large/zhoujinjuan_data/data/result-3/car_' + name + '.png', frame_copy)
-        # elif all_face_embeddings and all_licenses:
-        #     name = str(uuid.uuid4())[:7]
-        #     cv2.imwrite('/mnt/large/zhoujinjuan_data/data/result-3/all_' + name + '.png', frame_copy)
-        return all_face_embeddings, all_licenses
-
-
-if __name__ == '__main__':
-    import os
-    obj = TrackMain()
-    video_path = os.path.join('/mnt/large/zhoujinjuan_data/data/4.mp4')
-    cap = cv2.VideoCapture(video_path)
-    all_face_embeds = []
-    all_licenses_list = []
-    frame_count = 0
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame_count += 1
-        face_embeds, license_list = obj.process_one_frame(frame)
-        all_face_embeds.extend(face_embeds)
-        all_licenses_list.extend(license_list)
-    print(f'frame count {frame_count}\n face embeddings {len(all_face_embeds)}\nlicenses {len(all_licenses_list)}')
+        return all_face_embeddings, all_licenses, frame_copy
 

src/main_v2_test.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Filter_Object 
+@File ：main_v2_test.py
+@IDE ：PyCharm 
+@Author ：zjj
+@Date ：2024/1/19 9:51 
+'''
+import os
+import time
+import unittest
+import uuid
+from pathlib import Path
+import cv2
+from main_v2 import TrackMain
+from utils import extract_yolo_results
+
+file_name = 'AM03778_09060120200408073818_0015A.MP4'
+
+
+def make_result_dir():
+    # 保存图片，可视化
+    # dir_name = str(uuid.uuid4())[:7]
+    dir_name = file_name
+    dirx = Path(r'/mnt/large/zhoujinjuan_data/data/result') / dir_name
+    if not dirx.exists():
+        dirx.mkdir()
+        print(f'tmp dir {dirx}')
+    return dirx
+
+
+class TestMain(unittest.TestCase):
+    def setUp(self):
+        self.use_trt = True
+        self.half = False
+        self.obj = TrackMain(device='gpu', use_trt=self.use_trt, half=self.half)
+        video_path = os.path.join('/mnt/large/zhoujinjuan_data/data/tp_videos/20240122-test/' + file_name)
+        self.cap = cv2.VideoCapture(video_path)
+        # 记录视频所有满足条件的人脸
+        self.all_face_embeds = []
+        # 记录视频所有满足条件的车牌号
+        self.all_licenses_list = []
+        self.frame_count = 0
+        # 跳帧：每x帧处理一帧
+        self.skip_frame = 1
+        print('init success')
+
+    def process_one_frame_v2(self, frame):
+        results = self.obj.yolo_seg_model(frame)
+        # 提取模型检测结果 infos=[xywhs, cls, scores, xy_list]
+        is_hit, infos = extract_yolo_results(results)
+        if not is_hit:
+            print('no target detect')
+            return [], [], frame
+
+        # yolo检测结果可视化
+        # annotated_frame = results[0].plot()
+        # cv2.imwrite('data/yolo_result.png', annotated_frame)
+
+        # 阈值筛选
+        tinfos_each_frame = self.obj.statistics_one_frame(infos)
+        # print(tinfos_each_frame)
+
+        # 保持原接口不动，比process_one_frame 多返回一个frame_copy
+        face_embeddings, licenses, frame_copy = self.obj.process_topn_in_one_frame(frame, tinfos_each_frame)
+        return face_embeddings, licenses, frame_copy
+
+    def test_process_one_frame(self):
+        """
+        测试一帧的检测结果
+        :return:
+        """
+        for index, name in enumerate(['truck2.png']):
+            pathx = 'data/' + name
+            frame = cv2.imread(pathx)
+            _, licenses, frame_copy = self.process_one_frame_v2(frame)
+            print(f'licenses {licenses}')
+            cv2.imwrite('data/result_' + str(index) + '.png', frame_copy)
+
+    def test_main(self):
+        print('begin test')
+        # 先创建结果目录
+        result_dir = make_result_dir()
+        start = time.time()
+        while self.cap.isOpened():
+            ret, frame = self.cap.read()
+            if not ret:
+                break
+            self.frame_count += 1
+
+            # todo 考虑跳帧
+            if self.frame_count % self.skip_frame != 0:
+                continue
+
+            # 每一帧的处理结果
+            face_embeds, license_list, frame_copy = self.process_one_frame_v2(frame)
+
+            # todo test 测试过程中保存图片可视化
+            if face_embeds and not license_list:
+                name = str(uuid.uuid4())[:7]
+                cv2.imwrite((result_dir / ('face_' + name + '.png')).as_posix(), frame_copy)
+            elif not face_embeds and license_list:
+                name = str(uuid.uuid4())[:7]
+                cv2.imwrite((result_dir / ('motorVehicle_' + name + '.png')).as_posix(), frame_copy)
+            elif face_embeds and license_list:
+                name = str(uuid.uuid4())[:7]
+                cv2.imwrite((result_dir / ('all_' + name + '.png')).as_posix(), frame_copy)
+
+            self.all_face_embeds.extend(face_embeds)
+            self.all_licenses_list.extend(license_list)
+        end = time.time()
+        t = round(end - start, 2)
+        actual_frame_number = round(self.frame_count / self.skip_frame)
+        print(f'frame count {self.frame_count}\nactual deal frame count {actual_frame_number}\ntime {t}s'
+              f'\nface embeddings count {len(self.all_face_embeds)}'
+              f'\nlicenses count {len(self.all_licenses_list)}')
+        print(self.all_licenses_list)
+        print(f'use_trt={self.use_trt}\nhalf={self.half}\n')
+
+
+if __name__ == '__main__':
+    unittest.main()

src/utils.py

@@ -27,7 +27,7 @@ from config import (
     FACE_COMPARE_THRESHOLD,
     ARTICULATION_MODEL_PATH,
     ARTICULATION_RANGE_PATH,
-    POLICE_IOU,
+    POLICE_IOU, DELTA_Y,
 )
 
 
@@ -497,12 +497,65 @@ def predict_ocr(frame):
     # OCR检测
     license_plate_list = []
     ocr_result = ocr_alg.ocr_predict(frame)
-    # 赋予空值
+    # 常规车牌号识别
     for txt in ocr_result[0].text:
         # 车牌
-        plate_num = parse_plate_number(txt)
-        if plate_num != None:
-            license_plate_list.append(plate_num)
+        if len(txt) in [7, 8, 9]:
+            plate_num = parse_plate_number(txt)
+            if plate_num is not None:
+                # 中间漏点的补上，通常是大卡车车厢上的车牌被识别
+                if len(txt) == 7:
+                    plate_num = plate_num[:2] + '·' + plate_num[2:]
+                license_plate_list.append(plate_num)
+
+    # todo 20240124 摩托车车牌规则匹配，通常是上下两行文本(不管上面是否检测到车牌)
+    moto_list = parse_moto_number(ocr_result)
+
+    license_plate_list.extend(moto_list)
+    return license_plate_list
+
+
+# 使用正则解析车牌号
+def parse_plate_number(txt):
+    pattern = "^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}[A-Z]{1}[·]?[A-Z0-9]{4}[A-Z0-9挂学警港澳]{1}$"
+    pattern1 = "^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}[A-Z]{1}[·]?[A-Z0-9]{5}[A-Z0-9挂学警港澳]{1}$"
+    match = re.match(pattern, txt)
+    match1 = re.match(pattern1, txt)
+    if match:
+        return match.group()
+    if match1:
+        return match1.group()
+
+
+def parse_moto_number(ocr_result):
+    """
+    摩托车车牌规则匹配，通常是上下两行文本
+    :param ocr_result:
+    :return:
+    """
+    license_plate_list = []
+    p_before = "^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}[·]?[A-Z]{1}$"
+    p_end = "^[A-Z0-9]{4}[A-Z0-9挂学警港澳]{1}$"
+    p_end1 = "^[A-Z0-9]{5}[A-Z0-9挂学警港澳]{1}$"
+    box_text_list = [[box, text] for box, text in zip(ocr_result[0].boxes, ocr_result[0].text)]
+    for index, box_text in enumerate(box_text_list):
+        box, text = box_text
+        # 文本长度预先过滤
+        if len(text) in [2, 3]:
+            match = re.match(p_before, text)
+            if match:
+                before = match.group()
+                if len(before) == 2:
+                    before = before[0] + '·' + before[1]
+                for box_j, text_j in box_text_list[index + 1:]:
+                    # 先用文本长度过滤，然后用上面box的ymax和下面box的ymin差距小于阈值筛选
+                    if len(text_j) in [5, 6] and (abs(box_j[1] - box[5]) < DELTA_Y or abs(box_j[3] - box[5]) < DELTA_Y):
+                        match_j = re.match(p_end, text_j)
+                        if not match_j:
+                            match_j = re.match(p_end1, text_j)
+                        if match_j:
+                            license_plate_list.append(before + match_j.group())
+
     return license_plate_list
 
 
@@ -573,18 +626,6 @@ def output_car_targets(left_cars, save_dir, output_car_dir):
     return paths
 
 
-# 使用正则解析车牌号
-def parse_plate_number(txt):
-    pattern = "^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}[A-Z]{1}[·]{1}[A-Z0-9]{4}[A-Z0-9挂学警港澳]{1}$"
-    pattern1 = "^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}[A-Z]{1}[·]{1}[A-Z0-9]{5}[A-Z0-9挂学警港澳]{1}$"
-    match = re.match(pattern, txt)
-    match1 = re.match(pattern1, txt)
-    if match:
-        return match.group()
-    if match1:
-        return match1.group()
-
-
 def det_articulation(image=None, img_path=None):
     if image is None and not img_path:
         return 0

src/utils_test.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Filter_Object 
+@File ：utils_test.py
+@IDE ：PyCharm 
+@Author ：zjj
+@Date ：2024/1/25 9:56 
+'''
+import unittest
+
+import cv2
+
+from algs import ocr_alg
+from utils import parse_moto_number, predict_ocr, parse_plate_number
+
+
+class UtilsTest(unittest.TestCase):
+
+    @staticmethod
+    def test_predict_ocr():
+        frame = cv2.imread('data/truck1.png')
+        result = predict_ocr(frame)
+        print(result)
+
+    @staticmethod
+    def test_parse_plate_number():
+        text = '苏J7272Z'
+        result = parse_plate_number(text)
+        print(result)
+
+    @staticmethod
+    def test_parse_moto_number():
+        frame = cv2.imread('data/moto2.jpg')
+        ocr_result = ocr_alg.ocr_predict(frame)
+        moto_list = parse_moto_number(ocr_result)
+        print(moto_list)