修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流

修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流!

前言:最近在做YOLO项目到时候,读取网络摄像头流的时候,由于网络的原因,总是发生断流的现象,导致程序奔溃,认真分析源码后,发现作者在开始读取摄像头的时候用的是LoadStreams,用于多路视频同时检测;但在utils包下,发现LoadWebcam类,是用于读取单路的摄像头,但启动不了,经过认知分析修改了一下,成功的启用LoadWebcam类来检测rtsp流,并为其添加了防止断流的机制,具体的教程如下:

1、进入datasets.py下,打开LoadWebcam,然后复制代码,新建一个Loadwebcam_web类:

  • YOLOv5 5.0 原作者的LoadWebcam类
class LoadWebcam:  # for inference
    def __init__(self, pipe='0', img_size=640, stride=32):
        self.img_size = img_size
        self.stride = stride

        if pipe.isnumeric():
            pipe = eval(pipe)  # local camera
        # pipe = 'rtsp://192.168.1.64/1'  # IP camera
        # pipe = 'rtsp://username:[email protected]/1'  # IP camera with login
        # pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg'  # IP golf camera

        self.pipe = pipe
        self.cap = cv2.VideoCapture(pipe)  # video capture object
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3)  # set buffer size

    def __iter__(self):
        self.count = -1
        return self

    def __next__(self):
        self.count += 1
        if cv2.waitKey(1) == ord('q'):  # q to quit
            self.cap.release()
            cv2.destroyAllWindows()
            raise StopIteration

        # Read frame
        if self.pipe == 0:  # local camera
            ret_val, img0 = self.cap.read()
            img0 = cv2.flip(img0, 1)  # flip left-right
        else:  # IP camera
            n = 0
            while True:
                n += 1
                self.cap.grab()
                if n % 30 == 0:  # skip frames
                    ret_val, img0 = self.cap.retrieve()
                    if ret_val:
                        break

        # Print
        assert ret_val, f'Camera Error {self.pipe}'
        img_path = 'webcam.jpg'
        print(f'webcam {self.count}: ', end='')

        # Padded resize
        img = letterbox(img0, self.img_size, stride=self.stride)[0]

        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)

        return img_path, img, img0, None

    def __len__(self):
        return 0
  • 自己新建的LoadWebcam_web类:
class LoadWebcam_web:  # for inference
    def __init__(self, pipe='0', img_size=640, stride=32):
        self.mode = 'video' # 配合5.0的loadwebcam新加的东西
        self.img_size = img_size
        self.stride = stride

        if pipe.isnumeric():
            pipe = eval(pipe)  # local camera
        # pipe = 'rtsp://192.168.1.64/1'  # IP camera
        # pipe = 'rtsp://username:[email protected]/1'  # IP camera with login
        # pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg'  # IP golf camera

        self.pipe = pipe
        self.cap = cv2.VideoCapture(pipe)  # video capture object
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3)  # set buffer size

    def __iter__(self):
        self.count = -1
        return self

    def __next__(self):
        self.count += 1
        if cv2.waitKey(1) == ord('q'):  # q to quit
            self.cap.release()
            cv2.destroyAllWindows()
            raise StopIteration

        # Read frame
        if self.pipe == 0:  # local camera
            ret_val, img0 = self.cap.read()
            img0 = cv2.flip(img0, 1)  # flip left-right
        else:  # IP camera
            n = 0
            while True:
                n += 1
                self.cap.grab()
                if n % 4 == 0:  # skip frames  抽帧检测,每1帧或者每4帧检测一次!
                    ret_val, img0 = self.cap.retrieve()
                    # ret_val 是控制图片的标志,图片如果正确的放回来的话,是true,如果没有获取图片则是false
                    if not ret_val:
                        # print("aaaaaaaaaaaa")
                        print("发生了断流")
                        self.cap.release()  # 发生断流以后,释放掉摄像头
                        self.cap = cv2.VideoCapture(self.pipe)  # 重新连接摄像头
                        ret_val, img0 = self.cap.retrieve()  # 重新读取视频帧
                        print("重新进行了连接")
                    else:
                        break

        # Print
        assert ret_val, f'Camera Error {self.pipe}'
        img_path = 'webcam.jpg'
        print(f'webcam {self.count}: ', end='')

        # Padded resize
        img = letterbox(img0, self.img_size, stride=self.stride)[0]

        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img)

        return img_path, img, img0, None

    def __len__(self):
        return 0
  • 对比LoadWebcam类和LoadWebcam_web类可以发现修改了2处:

    • 增加了一个mode属性,目的是配合dataset的返回值使用

修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流_第1张图片

  • 添加了断流机制

修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流_第2张图片

2、回到detect.py文件

  • 导入LoadWebcam_web 这个类
from utils.datasets import LoadStreams, LoadImages,LoadWebcam,LoadWebcam_web
  • 修改此处的代码

修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流_第3张图片

总结:

使用前需要在detect1.2.py文件中导入LoadWebcam_web类;并在Process detections的时候修改一下代码:
将 p, s, im0, frame = path[i], '%g: ’ % i, im0s[i].copy(), dataset.count
换成:(目的是为了保证LoadWebcam_web的正常运行,不然是花屏)
p, s, im0, frame = path, ‘’, im0s, getattr(dataset, ‘frame’, 0):

3、测试

  • 断开网络后,检测到断流

  • 修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流_第4张图片

  • 连接网络后,又开始重新检测

  • 修改YOLOv5 5.0源代码utils包下的datasets.py文件中的LoadWebcam类,防止读取网络摄像头RTSP流时发生断流_第5张图片

你可能感兴趣的:(网络,计算机视觉,python,深度学习)