From 36b0a9e7d01637a71ef2050d16d945de5dad7929 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 17 May 2021 12:27:40 +0200 Subject: [PATCH] Add CAP_PROP_FRAME_COUNT for YouTube sources (#3193) --- utils/datasets.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/utils/datasets.py b/utils/datasets.py index f380b20b97d4..56db9d2f13b5 100755 --- 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