Implementing Resource Release Mechanism for Stream Loading (ultralyti…

…cs#4285)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <[email protected]>

ultralytics/data/loaders.py

@@ -34,13 +34,15 @@ class LoadStreams:
     def __init__(self, sources='file.streams', imgsz=640, vid_stride=1):
         """Initialize instance variables and check for consistent input stream shapes."""
         torch.backends.cudnn.benchmark = True  # faster for fixed-size inference
+        self.running = True  # running flag for Thread
         self.mode = 'stream'
         self.imgsz = imgsz
         self.vid_stride = vid_stride  # video frame-rate stride
         sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
         n = len(sources)
         self.sources = [ops.clean_str(x) for x in sources]  # clean source names for later
         self.imgs, self.fps, self.frames, self.threads, self.shape = [[]] * n, [0] * n, [0] * n, [None] * n, [None] * n
+        self.caps = [None] * n  # video capture objects
         for i, s in enumerate(sources):  # index, source
             # Start thread to read frames from video stream
             st = f'{i + 1}/{n}: {s}... '
@@ -51,21 +53,22 @@ def __init__(self, sources='file.streams', imgsz=640, vid_stride=1):
             if s == 0 and (is_colab() or is_kaggle()):
                 raise NotImplementedError("'source=0' webcam not supported in Colab and Kaggle notebooks. "
                                           "Try running 'source=0' in a local environment.")
-            cap = cv2.VideoCapture(s)
-            if not cap.isOpened():
+            self.caps[i] = cv2.VideoCapture(s)  # store video capture object
+            if not self.caps[i].isOpened():
                 raise ConnectionError(f'{st}Failed to open {s}')
-            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            fps = cap.get(cv2.CAP_PROP_FPS)  # warning: may return 0 or nan
-            self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf')  # infinite stream fallback
+            w = int(self.caps[i].get(cv2.CAP_PROP_FRAME_WIDTH))
+            h = int(self.caps[i].get(cv2.CAP_PROP_FRAME_HEIGHT))
+            fps = self.caps[i].get(cv2.CAP_PROP_FPS)  # warning: may return 0 or nan
+            self.frames[i] = max(int(self.caps[i].get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float(
+                'inf')  # infinite stream fallback
             self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30  # 30 FPS fallback
 
-            success, im = cap.read()  # guarantee first frame
+            success, im = self.caps[i].read()  # guarantee first frame
             if not success or im is None:
                 raise ConnectionError(f'{st}Failed to read images from {s}')
             self.imgs[i].append(im)
             self.shape[i] = im.shape
-            self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True)
+            self.threads[i] = Thread(target=self.update, args=([i, self.caps[i], s]), daemon=True)
             LOGGER.info(f'{st}Success ✅ ({self.frames[i]} frames of shape {w}x{h} at {self.fps[i]:.2f} FPS)')
             self.threads[i].start()
         LOGGER.info('')  # newline
@@ -76,7 +79,7 @@ def __init__(self, sources='file.streams', imgsz=640, vid_stride=1):
     def update(self, i, cap, stream):
         """Read stream `i` frames in daemon thread."""
         n, f = 0, self.frames[i]  # frame number, frame array
-        while cap.isOpened() and n < f:
+        while self.running and cap.isOpened() and n < f:
             # Only read a new frame if the buffer is empty
             if not self.imgs[i]:
                 n += 1
@@ -92,6 +95,19 @@ def update(self, i, cap, stream):
             else:
                 time.sleep(0.01)  # wait until the buffer is empty
 
+    def close(self):
+        """Close stream loader and release resources."""
+        self.running = False  # stop flag for Thread
+        for i, thread in enumerate(self.threads):
+            if thread.is_alive():
+                thread.join(timeout=5)  # Add timeout
+        for cap in self.caps:  # Iterate through the stored VideoCapture objects
+            try:
+                cap.release()  # release video capture
+            except Exception as e:
+                LOGGER.warning(f'WARNING ⚠️ Could not release VideoCapture object: {e}')
+        cv2.destroyAllWindows()
+
     def __iter__(self):
         """Iterates through YOLO image feed and re-opens unresponsive streams."""
         self.count = -1