Fix LoadStreams() dataloader frame skip issue (#3833)

* Update datasets.py to read every 4th frame of streams

* Update datasets.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Feras OughaliandGlenn Jocher authored and GitHub committed 2021-06-30 12:11:29 +02:00
1 parent 57c5d02bbe
commit 7d6af69638
1 file changed
+7 -24
+7 -24
View File
@@ -4,7 +4,6 @@ import glob
import hashlib import hashlib
import json import json
import logging import logging
import math
import os import os
import random import random
import shutil import shutil
@@ -15,6 +14,7 @@ from pathlib import Path
from threading import Thread from threading import Thread
import cv2 import cv2
import math
import numpy as np import numpy as np
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
@@ -210,15 +210,8 @@ class LoadWebcam: # for inference
def __init__(self, pipe='0', img_size=640, stride=32): def __init__(self, pipe='0', img_size=640, stride=32):
self.img_size = img_size self.img_size = img_size
self.stride = stride self.stride = stride
self.pipe = eval(pipe) if pipe.isnumeric() else pipe
if pipe.isnumeric(): self.cap = cv2.VideoCapture(self.pipe) # video capture object
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 self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
def __iter__(self): def __iter__(self):
@@ -233,18 +226,8 @@ class LoadWebcam: # for inference
raise StopIteration raise StopIteration
# Read frame # Read frame
if self.pipe == 0: # local camera ret_val, img0 = self.cap.read()
ret_val, img0 = self.cap.read() img0 = cv2.flip(img0, 1) # flip left-right
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 # Print
assert ret_val, f'Camera Error {self.pipe}' assert ret_val, f'Camera Error {self.pipe}'
@@ -308,12 +291,12 @@ class LoadStreams: # multiple IP or RTSP cameras
def update(self, i, cap): def update(self, i, cap):
# Read stream `i` frames in daemon thread # Read stream `i` frames in daemon thread
n, f = 0, self.frames[i] n, f, read = 0, self.frames[i], 1 # frame number, frame array, inference every 'read' frame
while cap.isOpened() and n < f: while cap.isOpened() and n < f:
n += 1 n += 1
# _, self.imgs[index] = cap.read() # _, self.imgs[index] = cap.read()
cap.grab() cap.grab()
if n % 4: # read every 4th frame if n % read == 0:
success, im = cap.retrieve() success, im = cap.retrieve()
self.imgs[i] = im if success else self.imgs[i] * 0 self.imgs[i] = im if success else self.imgs[i] * 0
time.sleep(1 / self.fps[i]) # wait time time.sleep(1 / self.fps[i]) # wait time