快（50%）两步实现yolov5只检测指定区域，屏蔽非检测区域（包含效果视频）

博雅塔下的灯光

已于 2024-07-11 12:23:22 修改

阅读量2.1k

点赞数 8

文章标签： YOLO opencv 计算机视觉

于 2023-12-20 09:37:15 首次发布

本文链接：https://blog.csdn.net/weixin_41826859/article/details/135099716

版权

话不多说，直接上效果

# 结果

屏蔽前

屏蔽前

屏蔽后

屏蔽后

第一步：制作mask

制作mask的方法有很多种，题主见过用labelme和用鼠标做的，这里推荐用鼠标获取mask多边形的的坐标。

获取坐标点，这里推荐使用cv2的鼠标事件

import cv2
# 鼠标回调函数
def mouse_callback(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        print('Mouse position: ', x, y)
# 读取图像
img = cv2.imread('image.jpg')
# 创建一个窗口
cv2.namedWindow('image')
# 设置鼠标回调函数
cv2.setMouseCallback('image', mouse_callback)
while True:
    # 显示图像
    cv2.imshow('image', img)
    # 等待按键，如果按下'q'键，退出循环
    if cv2.waitKey(0) & 0xFF == ord('q'):
        break
cv2.destroyAllWindows()

获取mask

		# mask的多边形坐标点
        pts = np.array([[0, 0], [0, 1080], [1920, 1080], [1920, 850], [850, 0]], dtype=np.int32)
        # 重塑为形状 (1, -1, 2)
        pts = pts.reshape((-1, 1, 2))
        # 创建一个全黑的掩码
        mask = np.zeros(im0.shape, dtype=np.uint8)
        # 使用白色填充这个多边形
        cv2.fillPoly(mask, [pts], (255, 255, 255))

屏蔽区域如下图黑色部分
在这里插入图片描述

第二步：修改相应的dataloader

题主使用的是offline的视频，因此使用了dataloader中的LoadImages类，如果使用在线视频流，可以修改LoadStreams类相应位置

找到修改位置
修改位置一般在最后return 前的letterbox前，题主使用了一个flag来决定是否屏蔽其他区域
代码修改

# 原来的代码
if self.transforms:
     im = self.transforms(im0)  # transforms
else:
     im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
     im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
     im = np.ascontiguousarray(im)  # contiguous

#修改后的代码
flag = 0
if flag:
 	# 关建行，使用位与运算来将黑色区域屏蔽掉
     mask_im = cv2.bitwise_and(im0, mask)
     if self.transforms:
         im = self.transforms(mask_im)  # transforms
     else:
         im = letterbox(mask_im, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
         im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
         im = np.ascontiguousarray(im)  # contiguous
 else:
     if self.transforms:
         im = self.transforms(im0)  # transforms
     else:
         im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
         im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
         im = np.ascontiguousarray(im)  # contiguous