+
18
-

yolo如何添加视频遮罩只统计指定区域?

yolo如何添加视频遮罩只统计指定区域?

网友回复

+
3
-

在 YOLO 中添加视频遮罩来只统计指定区域

我们可以通过以下几种方法实现:

1. 使用 ROI (Region of Interest) 蒙版方法
import cv2
import numpy as np
from ultralytics import YOLO

# 加载YOLO模型
model = YOLO('yolov8n.pt')

# 创建蒙版
def create_mask(frame_shape, points):
    mask = np.zeros(frame_shape[:2], dtype=np.uint8)
    # 将点转换为numpy数组
    points = np.array(points, dtype=np.int32)
    # 填充多边形区域
    cv2.fillPoly(mask, [points], 255)
    return mask

# 视频处理
def process_video(video_path, roi_points):
    cap = cv2.VideoCapture(video_path)

    # 获取第一帧来创建蒙版
    ret, frame = cap.read()
    if not ret:
        return

    # 创建ROI蒙版
    mask = create_mask(frame.shape, roi_points)

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 应用蒙版
        masked_frame = cv2.bitwise_and(frame, frame, mask=mask)

        # YOLO检测
        results = model(masked_frame)

        # 处理检测结果
        for r in results:
            boxes = r.boxes
            for box in boxes:
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0]
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

                # 检查边界框是否在ROI内
                box_center = ((x1 + x2) // 2, (y1 + y2) // 2)
                if mask[box_center[1], box_center[0]] > 0:
                    # 绘制边界框
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

                    # 显示类别和置信度
                    conf = box.conf[0]
                    cls = box.cls[0]
                    label = f'{model.names[int(cls)]} {conf:.2f}'
                    cv2.putText(frame, label, (x1, y1 - 10), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # 绘制ROI区域
        cv2.polylines(frame, [roi_points], True, (0, 0, 255), 2)

        # 显示结果
        cv2.imshow('YOLO with ROI', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

# 使用示例
if __name__ == "__main__":
    # 定义ROI区域的点(根据需要修改坐标)
    roi_points = np.array([
        [100, 100],   # 左上
        [500, 100],   # ...

点击查看剩余70%

我知道答案,我要回答