diff --git a/algorithm/check.py b/algorithm/check.py new file mode 100644 index 0000000..957af3e --- /dev/null +++ b/algorithm/check.py @@ -0,0 +1,208 @@ +import cv2 +import supervision as sv +from ultralytics import YOLO +import os +import time +import numpy as np +from tqdm import tqdm +from PIL import Image, ImageDraw, ImageFont +import io + +def put_chinese_text(img, text, position, font_size=20, color=(0, 0, 255)): + # 将OpenCV图像转换为PIL图像 + img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) + draw = ImageDraw.Draw(img_pil) + + # 加载支持中文的字体 + try: + # 尝试加载系统中的微软雅黑字体 + font = ImageFont.truetype("msyh.ttc", font_size) + except: + # 如果找不到指定字体,使用默认字体 + font = ImageFont.load_default() + + # 绘制文本 + draw.text(position, text, font=font, fill=color) + + # 将PIL图像转换回OpenCV图像 + return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) + +def process_video_with_detection( + input_video_path: str, + output_video_path: str, + model_name: str = "yoll11n.pt", + confidence_threshold: float = 0.5, + classes: list = None, + show_live: bool = True, + save_annotated: bool = False +): + """ + 参数: + input_video_path: 输入视频文件路径 + output_video_path: 输出视频文件路径 + model_name: YOLO模型名称或路径 (yoll11n.pt) + confidence_threshold: 检测置信度阈值 (0-1) + classes: 要检测的类别ID列表 (None表示所有类别) + show_live: 是否实时显示处理过程 + save_annotated: 是否保存标注后的视频 + """ + # 1. 初始化模型 + print(f"加载模型: {model_name}") + model = YOLO(model_name) + + # 2. 初始化视频读取器 + print(f"打开视频文件: {input_video_path}") + video_info = sv.VideoInfo.from_video_path(input_video_path) + cap = cv2.VideoCapture(input_video_path) + + if not cap.isOpened(): + print(f"无法打开视频文件: {input_video_path}") + return + + # 3. 初始化视频写入器 (如果需要保存结果) + if save_annotated: + output_dir = os.path.dirname(output_video_path) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir) + print(f"创建输出目录: {output_dir}") + + writer = cv2.VideoWriter( + output_video_path, + cv2.VideoWriter_fourcc(*'mp4v'), + video_info.fps, + (video_info.width, video_info.height) + ) + print(f"准备保存结果到: {output_video_path}") + + # 4. 初始化Supervision工具 + byte_tracker = sv.ByteTrack() + box_annotator = sv.BoxAnnotator( + + ) + label_annotator = sv.LabelAnnotator() + trace_annotator = sv.TraceAnnotator() + + # 5. 处理进度跟踪 + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + print(f"视频总帧数: {total_frames}") + pbar = tqdm(total=total_frames, desc="处理视频帧") + + # 6. 处理每一帧 + frame_count = 0 + processing_times = [] + + while cap.isOpened(): + start_time = time.time() + + # 读取帧 + ret, frame = cap.read() + if not ret: + break + + frame_count += 1 + + # 使用YOLO进行目标检测 + results = model( + frame, + conf=confidence_threshold, + classes=classes, + verbose=False + )[0] + + # 转换为Supervision检测对象 + detections = sv.Detections.from_ultralytics(results) + + + # 准备标注 + labels = [ + f"{results.names[class_id]} {confidence:.2f}" + for class_id, confidence in + zip(detections.class_id, detections.confidence) + ] + + # 标注边界框 + annotated_frame = box_annotator.annotate( + scene=frame.copy(), + detections=detections + ) + + # 标注标签 + annotated_frame = label_annotator.annotate( + scene=annotated_frame, + detections=detections, + labels=labels + ) + + # 定义多边形点 + pts1 = [(1153.11, 273.86), (1146.09, 370.77), (1217.71, 352.51), (1220.52, 257.01)] + pts2 = [(1140.47, 506.99), (1214.90, 504.19), (1212.09, 592.66), (1136.25, 594.07)] + # 转换为整数坐标 + pts1 = [(int(x), int(y)) for x, y in pts1] + pts2 = [(int(x), int(y)) for x, y in pts2] + # 绘制多边形 + cv2.polylines(annotated_frame, [np.array(pts1)], isClosed=True, color=(0, 255, 0), thickness=2) + cv2.polylines(annotated_frame, [np.array(pts2)], isClosed=True, color=(0, 0, 255), thickness=2) + + # 检测目标是否通过多边形 + polygon = np.array(pts1, np.int32) + warning_displayed = False + + # 遍历所有检测到的目标 + for bbox in detections.xyxy: + # 获取目标边界框 + x1, y1, x2, y2 = bbox + # 计算中心点 + center_x = int((x1 + x2) / 2) + center_y = int((y1 + y2) / 2) + + # 检测中心点是否在多边形内 + distance = cv2.pointPolygonTest(polygon, (center_x, center_y), False) + + # 如果点在多边形内 (distance >= 0) + if distance >= 0: + # 显示警告信息 + annotated_frame = put_chinese_text(annotated_frame, "警告: 目标进入区域!", (50, 50), font_size=20, color=(255, 0, 0)) + + # 计算处理时间 + end_time = time.time() + processing_time = end_time - start_time + processing_times.append(processing_time) + + # 实时显示处理结果 + if show_live: + cv2.imshow("detect", annotated_frame) + if cv2.waitKey(1) == 27: # ESC键退出 + break + + # 保存处理后的帧 + if save_annotated: + writer.write(annotated_frame) + + # 更新进度条 + pbar.update(1) + + # 7. 清理资源 + cap.release() + if save_annotated: + writer.release() + if show_live: + cv2.destroyAllWindows() + pbar.close() + +if __name__ == "__main__": + # 示例用法 + input_video = r"D:\yolo8\mp4\12333.mp4" # 替换为您的视频路径 + output_video = r"output/annotated_video.mp4" # 输出视频路径 + + # 选择要检测的类别 + selected_classes = [0] + + process_video_with_detection( + input_video_path=input_video, + output_video_path=output_video, + model_name=r"D:\yolo8\结果文件\登轮\weights\best.pt", # 使用预训练模型 + confidence_threshold=0.5, # 置信度阈值 + classes=selected_classes, # 指定检测类别 + show_live=True, # 实时显示处理 + save_annotated=False # 保存结果视频 + ) \ No newline at end of file