Feras Oughali glenn-jocher commited on
Commit
7d6af69
1 Parent(s): 57c5d02

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

Browse files

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

* Update datasets.py

Co-authored-by: Glenn Jocher <[email protected]>

Files changed (1) hide show
  1. utils/datasets.py +7 -24
utils/datasets.py CHANGED
@@ -4,7 +4,6 @@ import glob
4
  import hashlib
5
  import json
6
  import logging
7
- import math
8
  import os
9
  import random
10
  import shutil
@@ -15,6 +14,7 @@ from pathlib import Path
15
  from threading import Thread
16
 
17
  import cv2
 
18
  import numpy as np
19
  import torch
20
  import torch.nn.functional as F
@@ -210,15 +210,8 @@ class LoadWebcam: # for inference
210
  def __init__(self, pipe='0', img_size=640, stride=32):
211
  self.img_size = img_size
212
  self.stride = stride
213
-
214
- if pipe.isnumeric():
215
- pipe = eval(pipe) # local camera
216
- # pipe = 'rtsp://192.168.1.64/1' # IP camera
217
- # pipe = 'rtsp://username:[email protected]/1' # IP camera with login
218
- # pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg' # IP golf camera
219
-
220
- self.pipe = pipe
221
- self.cap = cv2.VideoCapture(pipe) # video capture object
222
  self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
223
 
224
  def __iter__(self):
@@ -233,18 +226,8 @@ class LoadWebcam: # for inference
233
  raise StopIteration
234
 
235
  # Read frame
236
- if self.pipe == 0: # local camera
237
- ret_val, img0 = self.cap.read()
238
- img0 = cv2.flip(img0, 1) # flip left-right
239
- else: # IP camera
240
- n = 0
241
- while True:
242
- n += 1
243
- self.cap.grab()
244
- if n % 30 == 0: # skip frames
245
- ret_val, img0 = self.cap.retrieve()
246
- if ret_val:
247
- break
248
 
249
  # Print
250
  assert ret_val, f'Camera Error {self.pipe}'
@@ -308,12 +291,12 @@ class LoadStreams: # multiple IP or RTSP cameras
308
 
309
  def update(self, i, cap):
310
  # Read stream `i` frames in daemon thread
311
- n, f = 0, self.frames[i]
312
  while cap.isOpened() and n < f:
313
  n += 1
314
  # _, self.imgs[index] = cap.read()
315
  cap.grab()
316
- if n % 4: # read every 4th frame
317
  success, im = cap.retrieve()
318
  self.imgs[i] = im if success else self.imgs[i] * 0
319
  time.sleep(1 / self.fps[i]) # wait time
 
4
  import hashlib
5
  import json
6
  import logging
 
7
  import os
8
  import random
9
  import shutil
 
14
  from threading import Thread
15
 
16
  import cv2
17
+ import math
18
  import numpy as np
19
  import torch
20
  import torch.nn.functional as F
 
210
  def __init__(self, pipe='0', img_size=640, stride=32):
211
  self.img_size = img_size
212
  self.stride = stride
213
+ self.pipe = eval(pipe) if pipe.isnumeric() else pipe
214
+ self.cap = cv2.VideoCapture(self.pipe) # video capture object
 
 
 
 
 
 
 
215
  self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
216
 
217
  def __iter__(self):
 
226
  raise StopIteration
227
 
228
  # Read frame
229
+ ret_val, img0 = self.cap.read()
230
+ img0 = cv2.flip(img0, 1) # flip left-right
 
 
 
 
 
 
 
 
 
 
231
 
232
  # Print
233
  assert ret_val, f'Camera Error {self.pipe}'
 
291
 
292
  def update(self, i, cap):
293
  # Read stream `i` frames in daemon thread
294
+ n, f, read = 0, self.frames[i], 1 # frame number, frame array, inference every 'read' frame
295
  while cap.isOpened() and n < f:
296
  n += 1
297
  # _, self.imgs[index] = cap.read()
298
  cap.grab()
299
+ if n % read == 0:
300
  success, im = cap.retrieve()
301
  self.imgs[i] = im if success else self.imgs[i] * 0
302
  time.sleep(1 / self.fps[i]) # wait time