# -*- coding: utf-8 -*- # @Author : pan # @Description : 封装了一些函数 # @Date : 2023年8月8日09:45:23 import supervision as sv from ultralytics import YOLO from ultralytics.data.loaders import LoadStreams from ultralytics.engine.predictor import BasePredictor from ultralytics.utils import DEFAULT_CFG, SETTINGS from ultralytics.utils.torch_utils import smart_inference_mode from ultralytics.utils.files import increment_path from ultralytics.cfg import get_cfg from ultralytics.utils.checks import check_imshow from PySide6.QtCore import Signal, QObject from pathlib import Path import datetime import numpy as np import time import cv2 from classes.paint_trail import draw_trail from utils.main_utils import check_path x_axis_time_graph = [] y_axis_count_graph = [] video_id_count = 0 class YoloPredictor(BasePredictor, QObject): yolo2main_trail_img = Signal(np.ndarray) # 轨迹图像信号 yolo2main_box_img = Signal(np.ndarray) # 绘制了标签与锚框的图像的信号 yolo2main_status_msg = Signal(str) # 检测/暂停/停止/测试完成等信号 yolo2main_fps = Signal(str) # fps yolo2main_labels = Signal(dict) # 检测到的目标结果（每个类别的数量） yolo2main_progress = Signal(int) # 进度条 yolo2main_class_num = Signal(int) # 当前帧类别数 yolo2main_target_num = Signal(int) # 当前帧目标数 def __init__(self, cfg=DEFAULT_CFG, overrides=None): super(YoloPredictor, self).__init__() QObject.__init__(self) try: self.args = get_cfg(cfg, overrides) except: pass project = self.args.project or Path(SETTINGS['runs_dir']) / self.args.task name = f'{self.args.mode}' self.save_dir = increment_path(Path(project) / name, exist_ok=self.args.exist_ok) self.done_warmup = False if self.args.show: self.args.show = check_imshow(warn=True) # GUI args self.used_model_name = None # 使用过的检测模型名称 self.new_model_name = None # 新更改的模型 self.source = '' # 输入源str self.progress_value = 0 # 进度条的值 self.stop_dtc = False # 终止bool self.continue_dtc = True # 暂停bool # config self.iou_thres = 0.45 # iou self.conf_thres = 0.25 # conf self.speed_thres = 0.01 # delay, ms （缓冲） self.save_res = False # 保存MP4 self.save_txt = False # 保存txt self.save_res_path = "pre_result" self.save_txt_path = "pre_labels" self.show_labels = True # 显示图像标签bool self.show_trace = True # 显示图像轨迹bool # 运行时候的参数放这里 self.start_time = None # 拿来算FPS的计数变量 self.count = None self.sum_of_count = None self.class_num = None self.total_frames = None self.lock_id = None # 设置线条样式 厚度 & 缩放大小 self.box_annotator = sv.BoxAnnotator( thickness=2, text_thickness=1, text_scale=0.5 ) # 点击开始检测按钮后的检测事件 @smart_inference_mode() # 一个修饰器，用来开启检测模式：如果torch>=1.9.0，则执行torch.inference_mode()，否则执行torch.no_grad() def run(self): self.yolo2main_status_msg.emit('正在加载模型...') LoadStreams.capture = '' self.count = 0 # 拿来参与算FPS的计数变量 self.start_time = time.time() # 拿来算FPS的计数变量 global video_id_count # 检查保存路径 if self.save_txt: check_path(self.save_txt_path) if self.save_res: check_path(self.save_res_path) model = self.load_yolo_model() # 获取数据源 （不同的类型获取不同的数据源） iter_model = iter( model.track(source=self.source, show=False, stream=True, iou=self.iou_thres, conf=self.conf_thres)) # 折线图数据初始化 global x_axis_time_graph, y_axis_count_graph x_axis_time_graph = [] y_axis_count_graph = [] self.yolo2main_status_msg.emit('检测中...') # 使用OpenCV读取视频——获取进度条 if 'mp4' in self.source or 'avi' in self.source or 'mkv' in self.source or 'flv' in self.source or 'mov' in self.source: cap = cv2.VideoCapture(self.source) self.total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) cap.release() # 如果保存，则创建写入对象 img_res, result, height, width = self.recognize_res(iter_model) fourcc = cv2.VideoWriter_fourcc(*'XVID') out = None # 视频写出变量 if self.save_res: out = cv2.VideoWriter(f'{self.save_res_path}/video_result_{video_id_count}.mp4', fourcc, 25, (width, height), True) # 保存检测视频的路径 # 开始死循环检测 while True: try: # 暂停与开始 if self.continue_dtc: img_res, result, height, width = self.recognize_res(iter_model) self.res_address(img_res, result, height, width, model, out) # 终止 if self.stop_dtc: if self.save_res: if out: out.release() video_id_count += 1 self.source = None self.yolo2main_status_msg.emit('检测终止') self.release_capture() # 这里是为了终止使用摄像头检测函数的线程，改了yolo源码 break # 检测截止（本地文件检测） except StopIteration: if self.save_res: out.release() video_id_count += 1 print('writing complete') self.yolo2main_status_msg.emit('检测完成') self.yolo2main_progress.emit(1000) cv2.destroyAllWindows() # 单目标追踪停止！ self.source = None break try: out.release() except: pass # 进行识别——并返回所有结果 def res_address(self, img_res, result, height, width, model, out): # 复制一份 img_box = np.copy(img_res) # 右边的图（会绘制标签！） img_res是原图-不会受影响 img_trail = np.copy(img_res) # 左边的图 # 如果没有识别的： if result.boxes.id is None: # 目标都是0 self.sum_of_count = 0 self.class_num = 0 labels_write = "暂未识别到目标！" # 如果有识别的 else: detections = sv.Detections.from_yolov8(result) detections.tracker_id = result.boxes.id.cpu().numpy().astype(int) # id 、位置、目标总数 self.class_num = self.get_class_number(detections) # 类别数 id = detections.tracker_id # id xyxy = detections.xyxy # 位置 self.sum_of_count = len(id) # 目标总数 # 轨迹绘制部分 @@@@@@@@@@@@@@@@@@@@@@@@@@@@ if self.show_trace: img_trail = np.zeros((height, width, 3), dtype='uint8') # 黑布 identities = id grid_color = (255, 255, 255) line_width = 1 grid_size = 100 for y in range(0, height, grid_size): cv2.line(img_trail, (0, y), (width, y), grid_color, line_width) for x in range(0, width, grid_size): cv2.line(img_trail, (x, 0), (x, height), grid_color, line_width) draw_trail(img_trail, xyxy, model.model.names, id, identities) else: img_trail = img_res # 显示原图 # 画标签到图像上（并返回要写下的信息 labels_write, img_box = self.creat_labels(detections, img_box , model) # 写入txt——存储labels里的信息 if self.save_txt: with open(f'{self.save_txt_path}/result.txt', 'a') as f: f.write('当前时刻屏幕信息:' + str(labels_write) + f'检测时间: {datetime.datetime.now().strftime("%Y-%m-%d-%H:%M:%S")}' + f' 路段通过的目标总数: {self.sum_of_count}') f.write('\n') # 预测视频写入本地 if self.save_res: out.write(img_box) # 添加 折线图数据 now = datetime.datetime.now() new_time = now.strftime("%Y-%m-%d %H:%M:%S") if new_time not in x_axis_time_graph: # 防止同一秒写入 x_axis_time_graph.append(new_time) y_axis_count_graph.append(self.sum_of_count) # 抠锚框里的图 （单目标追踪） if self.lock_id is not None: self.lock_id = int(self.lock_id) self.open_target_tracking(detections=detections, img_res=img_res) # 传递信号给主窗口 self.emit_res(img_trail, img_box) # 识别结果处理 def recognize_res(self, iter_model): # 检测 ---然后获取有用的数据 result = next(iter_model) # 这里是检测的核心，每次循环都会检测一帧图像,可以自行打印result看看里面有哪些key可以用 img_res = result.orig_img # 原图 height, width, _ = img_res.shape return img_res, result, height, width # 单目标检测窗口开启 def open_target_tracking(self, detections, img_res): try: # 单目标追踪 ！！！！！ result_cropped = self.single_object_tracking(detections, img_res) # print(result_cropped) cv2.imshow(f'OBJECT-ID:{self.lock_id}', result_cropped) cv2.moveWindow(f'OBJECT-ID:{self.lock_id}', 0, 0) # press esc to quit if cv2.waitKey(5) & 0xFF == 27: self.lock_id = None cv2.destroyAllWindows() except: cv2.destroyAllWindows() pass # 单目标跟踪 def single_object_tracking(self, detections, img_box): store_xyxy_for_id = {} for xyxy, id in zip(detections.xyxy, detections.tracker_id): store_xyxy_for_id[id] = xyxy mask = np.zeros_like(img_box) try: if self.lock_id not in detections.tracker_id: cv2.destroyAllWindows() self.lock_id = None x1, y1, x2, y2 = int(store_xyxy_for_id[self.lock_id][0]), int(store_xyxy_for_id[self.lock_id][1]), int( store_xyxy_for_id[self.lock_id][2]), int(store_xyxy_for_id[self.lock_id][3]) cv2.rectangle(mask, (x1, y1), (x2, y2), (255, 255, 255), -1) result_mask = cv2.bitwise_and(img_box, mask) result_cropped = result_mask[y1:y2, x1:x2] result_cropped = cv2.resize(result_cropped, (256, 256)) return result_cropped except: cv2.destroyAllWindows() pass # 信号发送区 def emit_res(self, img_trail, img_box): time.sleep(self.speed_thres/1000) # 缓冲 # 轨迹图像（左边） self.yolo2main_trail_img.emit(img_trail) # 标签图（右边） self.yolo2main_box_img.emit(img_box) # 总类别数量 、 总目标数 self.yolo2main_class_num.emit(self.class_num) self.yolo2main_target_num.emit(self.sum_of_count) # 进度条 if '0' in self.source or 'rtsp' in self.source: self.yolo2main_progress.emit(0) else: self.progress_value = int(self.count / self.total_frames * 1000) self.yolo2main_progress.emit(self.progress_value) # 计算FPS self.count += 1 if self.count % 3 == 0 and self.count >= 3: # 计算FPS self.yolo2main_fps.emit(str(int(3 / (time.time() - self.start_time)))) self.start_time = time.time() # 加载模型 def load_yolo_model(self): if self.used_model_name != self.new_model_name: self.setup_model(self.new_model_name) self.used_model_name = self.new_model_name return YOLO(self.new_model_name) # 画标签到图像上 def creat_labels(self, detections, img_box, model): # 要画出来的信息 labels_draw = [ f"ID: {tracker_id} {model.model.names[class_id]}" for _, _, confidence, class_id, tracker_id in detections ] ''' 如果Torch装的是cuda版本的话：302行的代码需改成： labels_draw = [ f"OBJECT-ID: {tracker_id} CLASS: {model.model.names[class_id]} CF: {confidence:0.2f}" for _,confidence,class_id,tracker_id in detections ] ''' # 存储labels里的信息 labels_write = [ f"目标ID: {tracker_id} 目标类别: {class_id} 置信度: {confidence:0.2f}" for _, _, confidence, class_id, tracker_id in detections ] ''' 如果Torch装的是cuda版本的话：314行的代码需改成： labels_write = [ f"OBJECT-ID: {tracker_id} CLASS: {model.model.names[class_id]} CF: {confidence:0.2f}" for _,confidence,class_id,tracker_id in detections ] ''' # 如果显示标签 （要有才可以画呀！）---否则就是原图 if (self.show_labels == True) and (self.class_num != 0): img_box = self.box_annotator.annotate(scene=img_box, detections=detections, labels=labels_draw) return labels_write, img_box # 获取类别数 def get_class_number(self, detections): class_num_arr = [] for each in detections.class_id: if each not in class_num_arr: class_num_arr.append(each) return len(class_num_arr) # 释放摄像头 def release_capture(self): LoadStreams.capture = 'release' # 这里是为了终止使用摄像头检测函数的线程，改了yolo源码