修改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:password@192.168.1.64/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:password@192.168.1.64/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的返回值使用

• 添加了断流机制

2、回到detect.py文件

• 导入LoadWebcam_web 这个类

from utils.datasets import LoadStreams, LoadImages,LoadWebcam,LoadWebcam_web

• 修改此处的代码

总结：

使用前需要在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、测试

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

• 

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

•