|
- #!/usr/bin/env python3
- # -*- coding: utf-8 -*-
- # @Created : 2021/05/8
- # @Author : Koala
- # @FileName: vehicle.py
-
- import os
- import cv2
- import numpy as np
- from typing import Optional
- from run import TensorRTRun
- from utils import IoU
-
-
- def report_alert(t, x1, y1, x2, y2):
- print(f"There is a alert: {t, x1, y1, x2, y2}")
-
-
- class IllegalParking:
- def __init__(self, onnx_path='./sources/vehicl_model.onnx',
- trt_path='./sources/vehicl_model.trt',
- alert_boxes_filename='./sources/vehicl_alert_boxes.txt',
- max_batch_size=1, #支持多batch inference
- iou_threshold=0.3, #目标和警戒区域的重叠占比
- score_threshold=0.4, #目标检测器阈值
- detect_fps=25, #每秒检测的帧数
- alter_time=20, #连续检测时间,单位为秒
- alter_ratio=0.5, #检测丢帧比列
- same_threshold=0.7, #定位同一个目标的iou
- verbose=True, #是否可视化视频
- ):
- self.boxes = np.zeros([0, 4])
- self.runner = None # type: Optional[TensorRTRun]
- self.iou_threshold = iou_threshold
- self.score_threshold = score_threshold
- self.detect_fps = detect_fps # 每秒检测几帧的图片
- assert 0 <= alter_ratio <= 1
- self.alter_length = int(alter_time * detect_fps) # 停车被检测多少次后报警
- self.alter_number = int(self.alter_length * alter_ratio) # 相同车辆在最近<alter_length>被检测多少次被认为违停
- self.same_threshold = same_threshold # 被认为是同一辆车的iou阈值
- print(f'每秒检测: {self.detect_fps} 帧视频, 当IoU>={self.same_threshold}时, 认为是同一辆车')
- print(f'在连续的 {self.alter_length} 次检测中, 同一辆车至少检测到 {self.alter_number} 次则认为是违停并报警')
-
- self.read_alert_boxes(alert_boxes_filename)
- self.load_model(trt_path, onnx_path, max_batch_size, True)
- self._old_detected = np.zeros([0, 5], dtype=np.int)
- self._num_detected = []
- self.verbose = verbose
-
- ### 读取警戒框
- def read_alert_boxes(self, filename):
- assert os.path.isfile(filename), f'警戒框文件"{filename}"并不存在'
- with open(filename, 'r') as f:
- self.boxes = np.array([list(map(int, s.split())) for s in f.readlines() if s.strip()])
- print(f'从文件 {filename} 读取到的警戒框有:')
- for box in self.boxes:
- print(f'\t({box[0]}, {box[1]}), ({box[2]} {box[3]})')
-
- ### 加载或构建trt模型
- def load_model(self, trt_path, onnx_path='', max_batch_size=1, fp16=True):
- self.runner = TensorRTRun(trt_path)
- if os.path.isfile(trt_path):
- print(f'从"{trt_path}"加载trt模型')
- self.runner.load_engine()
- else:
- assert os.path.isfile(onnx_path), f'ONNX模型文件"{onnx_path}"并不存在'
- self.runner.build_engine(onnx_path, fp16=True, max_batch_size=max_batch_size)
- print(f'使用ONNX模型"{onnx_path}"构建trt模型, 并保存在"{trt_path}')
- self.runner.prepare()
- print('已完成准备工作')
- return
-
- def detect(self, img, alert_class):
- detected = self.runner(img)[0]
- if len(detected) == 0:
- return detected[:, :4]
-
- if alert_class == 0:
- detected = detected[detected[:, 5] == 0, :]
- print('机动车违规停车:')
- elif alert_class == 1:
- detected = detected[detected[:, 5] == 1, :]
- print('非机动车违规停车:')
- else:
- print('机动车或非机动车违规停车:')
- detected = detected[detected[:, 4] >= self.score_threshold, :4].astype(np.int)
-
- iou = IoU(detected, self.boxes)
- is_in_alert = np.any(iou > self.iou_threshold, axis=1)
- if self.verbose:
- for i, (x1, y1, x2, y2) in enumerate(detected):
- cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 1 + is_in_alert[i] * 2)
- return detected[is_in_alert, :]
-
- def __call__(self, video_path, alert_class):
- self._old_detected = np.zeros([0, 5], dtype=np.int)
- self._num_detected = []
-
- # assert os.path.isfile(video_path)
- if self.verbose:
- cv2.namedWindow(video_path)
- cap = cv2.VideoCapture(video_path)
- W = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
- H = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
- print(f'视频大小: WxH={W}x{H}')
- last_t = -1
- last_alert = []
-
- while True:
- ret, frame = cap.read()
- if not ret:
- return
- second = cap.get(cv2.CAP_PROP_POS_MSEC) * 1e-3
-
- now_t = int(second * self.detect_fps)
- if now_t != last_t:
- boxes = self.detect(frame, alert_class)
- n = len(boxes)
- if len(self._num_detected) == self.alter_length:
- self._old_detected = self._old_detected[self._num_detected[0]:, :]
- self._num_detected = self._num_detected[1:]
- boxes = np.concatenate([boxes, np.full([n, 1], now_t)], axis=1)
- last_alert = []
- if n > 0:
- is_same = IoU(boxes, self._old_detected, False) >= self.same_threshold
- for i in range(n):
- need_alert = np.unique(self._old_detected[is_same[i], 4]).size + 1 >= self.alter_number
- if need_alert:
- report_alert(second, *boxes[i, :4])
- last_alert.append(boxes[i, :4])
- self._num_detected.append(n)
- self._old_detected = np.concatenate([self._old_detected, boxes], axis=0)
- last_t = now_t
- if not self.verbose:
- continue
- ## 展示警戒框
- cv2.putText(frame, f'time: {second:.2f}s', (20, 50), 2, 2, (255, 255, 0))
- for x1, y1, x2, y2 in self.boxes:
- cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
- ## 展示违停车辆
- for x1, y1, x2, y2 in last_alert:
- # cv2.putText(frame, f'Illegal Parking', (x1, y1), 2, 2, (0, 255, 0))
- cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
- cv2.imshow(video_path, frame)
- k = cv2.waitKey(1)
- if k in [27, ord('q'), ord('Q')]:
- break
- elif k in [ord(' ')]:
- while True:
- cv2.imshow(video_path, frame)
- k = cv2.waitKey(100)
- if k in [27, ord('q'), ord('Q')]:
- return
- elif k == 32:
- break
- # break
- return
-
-
- def main(
- video_path ='1.mp4', #视频源
- alert_class = 2, #报警目标类型,0-机动车,1-非机动车,2-机动车和非机动车
- onnx_path ='./sources/vehicl_model.onnx', #onnx模型路径
- trt_path ='./sources/vehicl_model.trt', #trt序列化模型保存路径
- alert_boxes_filename ='./sources/vehicl_alert_boxes.txt', #违停区域坐标存储文件路径
- ):
- e = IllegalParking(onnx_path, trt_path, alert_boxes_filename)
- e(video_path, alert_class)
-
-
- if __name__ == '__main__':
- main()
|