diff --git a/ultralytics/data/loaders.py b/ultralytics/data/loaders.py index a59e8c8..88246e5 100644 --- a/ultralytics/data/loaders.py +++ b/ultralytics/data/loaders.py @@ -34,6 +34,7 @@ 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 @@ -41,6 +42,7 @@ class LoadStreams: 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 @@ class LoadStreams: 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 @@ class LoadStreams: 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 @@ class LoadStreams: 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