glenn-jocher commited on
Commit
36b0a9e
1 Parent(s): ae53f50

Add CAP_PROP_FRAME_COUNT for YouTube sources (#3193)

Browse files
Files changed (1) hide show
  1. utils/datasets.py +4 -3
utils/datasets.py CHANGED
@@ -172,7 +172,7 @@ class LoadImages: # for inference
172
  ret_val, img0 = self.cap.read()
173
 
174
  self.frame += 1
175
- print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
176
 
177
  else:
178
  # Read image
@@ -193,7 +193,7 @@ class LoadImages: # for inference
193
  def new_video(self, path):
194
  self.frame = 0
195
  self.cap = cv2.VideoCapture(path)
196
- self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
197
 
198
  def __len__(self):
199
  return self.nf # number of files
@@ -285,10 +285,11 @@ class LoadStreams: # multiple IP or RTSP cameras
285
  w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
286
  h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
287
  self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
 
288
 
289
  _, self.imgs[i] = cap.read() # guarantee first frame
290
  thread = Thread(target=self.update, args=([i, cap]), daemon=True)
291
- print(f' success ({w}x{h} at {self.fps:.2f} FPS).')
292
  thread.start()
293
  print('') # newline
294
 
 
172
  ret_val, img0 = self.cap.read()
173
 
174
  self.frame += 1
175
+ print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ', end='')
176
 
177
  else:
178
  # Read image
 
193
  def new_video(self, path):
194
  self.frame = 0
195
  self.cap = cv2.VideoCapture(path)
196
+ self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
197
 
198
  def __len__(self):
199
  return self.nf # number of files
 
285
  w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
286
  h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
287
  self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
288
+ self.frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
289
 
290
  _, self.imgs[i] = cap.read() # guarantee first frame
291
  thread = Thread(target=self.update, args=([i, cap]), daemon=True)
292
+ print(f" success ({f'{self.frames} frames ' if self.frames else ''}{w}x{h} at {self.fps:.2f} FPS).")
293
  thread.start()
294
  print('') # newline
295