Merge branch 'main' of http://36.103.203.89:3000/gezhonglin/border_Inspection

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        # 保存结果视频
+    )