From eb9fb02ca8d2257590f50612fff9cf67df41ce9a Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 21 May 2021 16:39:01 +0200
Subject: [PATCH 1/2] YouTube stream ending fix

Properly terminates YouTube streams on video end. Should resolve issues #2769 and #3220.
---
 utils/datasets.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/utils/datasets.py b/utils/datasets.py
index a14099875455..1dbed5e32a36 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -270,7 +270,7 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
             sources = [sources]
 
         n = len(sources)
-        self.imgs = [None] * n
+        self.imgs, self.fps, self.frames, self.stopped = [None] * n, [0] * n, [0] * n, [False] * n
         self.sources = [clean_str(x) for x in sources]  # clean source names for later
         for i, s in enumerate(sources):  # index, source
             # Start thread to read frames from video stream
@@ -284,12 +284,12 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
             assert cap.isOpened(), f'Failed to open {s}'
             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) or 30.0  # assume 30 FPS if cap gets 0 FPS
-            self.frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            self.fps[i] = (cap.get(cv2.CAP_PROP_FPS) % 100) or 30.0  # assume 30 FPS if cap gets 0 FPS
+            self.frames[i] = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or float('inf')  # assume infinite stream if 0 len
 
             _, self.imgs[i] = cap.read()  # guarantee first frame
             thread = Thread(target=self.update, args=([i, cap]), daemon=True)
-            print(f" success ({f'{self.frames} frames ' if self.frames else ''}{w}x{h} at {self.fps:.2f} FPS).")
+            print(f" success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
             thread.start()
         print('')  # newline
 
@@ -299,18 +299,18 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
         if not self.rect:
             print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
 
-    def update(self, index, cap):
-        # Read next stream frame in a daemon thread
-        n = 0
-        while cap.isOpened():
+    def update(self, i, cap):
+        # Read stream `i` frames in daemon thread
+        n, f = 0, self.frames[i]
+        while cap.isOpened() and n < f:
             n += 1
             # _, self.imgs[index] = cap.read()
             cap.grab()
-            if n == 4:  # read every 4th frame
+            if n % 4:  # read every 4th frame
                 success, im = cap.retrieve()
-                self.imgs[index] = im if success else self.imgs[index] * 0
-                n = 0
-            time.sleep(1 / self.fps)  # wait time
+                self.imgs[i] = im if success else self.imgs[i] * 0
+            time.sleep(1 / self.fps[i])  # wait time
+        self.stopped[i] = True
 
     def __iter__(self):
         self.count = -1
@@ -318,12 +318,12 @@ def __iter__(self):
 
     def __next__(self):
         self.count += 1
-        img0 = self.imgs.copy()
-        if cv2.waitKey(1) == ord('q'):  # q to quit
+        if all(self.stopped) or cv2.waitKey(1) == ord('q'):  # q to quit
             cv2.destroyAllWindows()
             raise StopIteration
 
         # Letterbox
+        img0 = self.imgs.copy()
         img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]
 
         # Stack

From 354b89dfcc9dcae77c9ef6cbd1e233805a523368 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 21 May 2021 16:47:12 +0200
Subject: [PATCH 2/2] Update datasets.py

---
 utils/datasets.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/utils/datasets.py b/utils/datasets.py
index 1dbed5e32a36..b05763386dc2 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -270,7 +270,7 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
             sources = [sources]
 
         n = len(sources)
-        self.imgs, self.fps, self.frames, self.stopped = [None] * n, [0] * n, [0] * n, [False] * n
+        self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
         self.sources = [clean_str(x) for x in sources]  # clean source names for later
         for i, s in enumerate(sources):  # index, source
             # Start thread to read frames from video stream
@@ -288,9 +288,9 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32):
             self.frames[i] = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or float('inf')  # assume infinite stream if 0 len
 
             _, self.imgs[i] = cap.read()  # guarantee first frame
-            thread = Thread(target=self.update, args=([i, cap]), daemon=True)
+            self.threads[i] = Thread(target=self.update, args=([i, cap]), daemon=True)
             print(f" success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
-            thread.start()
+            self.threads[i].start()
         print('')  # newline
 
         # check for common shapes
@@ -310,7 +310,6 @@ def update(self, i, cap):
                 success, im = cap.retrieve()
                 self.imgs[i] = im if success else self.imgs[i] * 0
             time.sleep(1 / self.fps[i])  # wait time
-        self.stopped[i] = True
 
     def __iter__(self):
         self.count = -1
@@ -318,7 +317,7 @@ def __iter__(self):
 
     def __next__(self):
         self.count += 1
-        if all(self.stopped) or cv2.waitKey(1) == ord('q'):  # q to quit
+        if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'):  # q to quit
             cv2.destroyAllWindows()
             raise StopIteration