Add CAP_PROP_FRAME_COUNT for YouTube sources (#3193)

ultralytics · May 17, 2021 · 36b0a9e · 36b0a9e
1 parent ae53f50
commit 36b0a9e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -172,7 +172,7 @@ def __next__(self):
                     ret_val, img0 = self.cap.read()
 
             self.frame += 1
-            print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
+            print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ', end='')
 
         else:
             # Read image
@@ -193,7 +193,7 @@ def __next__(self):
     def new_video(self, path):
         self.frame = 0
         self.cap = cv2.VideoCapture(path)
-        self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
     def __len__(self):
         return self.nf  # number of files
@@ -285,10 +285,11 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
             w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
             self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
+            self.frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
             _, self.imgs[i] = cap.read()  # guarantee first frame
             thread = Thread(target=self.update, args=([i, cap]), daemon=True)
-            print(f' success ({w}x{h} at {self.fps:.2f} FPS).')
+            print(f" success ({f'{self.frames} frames ' if self.frames else ''}{w}x{h} at {self.fps:.2f} FPS).")
             thread.start()
         print('')  # newline