修改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的返回值使用
- 添加了断流机制
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、测试
-
断开网络后，检测到断流
-
-
连接网络后，又开始重新检测
-
