add_0805

algorithm/check_seg.py

@@ -0,0 +1,169 @@
+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
+
+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_segmentation(
+    input_video_path: str,
+    output_video_path: str,
+    model_name: str = "yolov8n-seg.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分割模型名称或路径 (默认: yolov8n-seg.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工具 - 使用MaskAnnotator替代BoxAnnotator
+    mask_annotator = sv.MaskAnnotator(
+        color=sv.Color(r=0, g=255, b=0),  # 绿色掩码
+        opacity=0.5  # 50%透明度
+    )
+    
+    # 5. 处理进度跟踪
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    print(f"视频总帧数: {total_frames}")
+    pbar = tqdm(total=total_frames, desc="处理视频帧")
+    
+    
+    
+    # 7. 处理每一帧
+    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)
+        
+        # 应用分割掩码
+        segmented_frame = mask_annotator.annotate(
+            scene=frame.copy(),
+            detections=detections
+        )
+    
+        
+        # 计算处理时间
+        end_time = time.time()
+        processing_time = end_time - start_time
+        processing_times.append(processing_time)
+        
+        # 实时显示处理结果
+        if show_live:
+            cv2.imshow("Segmentation", segmented_frame)
+            if cv2.waitKey(1) == 27:  # ESC键退出
+                break
+        
+        # 保存处理后的帧
+        if save_annotated:
+            writer.write(segmented_frame)
+        
+        # 更新进度条
+        pbar.update(1)
+    
+    # 8. 清理资源
+    cap.release()
+    if save_annotated:
+        writer.release()
+    if show_live:
+        cv2.destroyAllWindows()
+    pbar.close()
+    
+    # 打印性能统计
+    if processing_times:
+        avg_time = sum(processing_times) / len(processing_times)
+        fps = 1 / avg_time if avg_time > 0 else 0
+        print(f"\n处理完成! 平均处理时间: {avg_time:.4f}秒/帧, 约 {fps:.2f} FPS")
+
+if __name__ == "__main__":
+    # 示例用法
+    input_video = r"D:\yolo8\新建文件夹\靠泊视频.mp4"  # 替换为您的视频路径
+    output_video = r"output/segmented_video.mp4"   # 输出视频路径
+    
+    # 选择要检测的类别 (0=人, 2=车等)
+    selected_classes = [0]  # 只检测人
+    
+    process_video_with_segmentation(
+        input_video_path=input_video,
+        output_video_path=output_video,
+        model_name=r"D:\yolo8\model\kaobo\weights\best.pt",  # 使用分割模型
+        confidence_threshold=0.3,  # 置信度阈值
+        classes=selected_classes,   # 指定检测类别
+        show_live=True,             # 实时显示处理
+        save_annotated=True         # 保存结果视频
+    )