✨ [Add] A stream dataloader for webcam

requirements.txt

@@ -3,6 +3,7 @@ graphviz
 hydra-core
 loguru
 numpy
+opencv-python
 Pillow
 pytest
 pyyaml

yolo/config/config.py

@@ -113,6 +113,7 @@ class NMSConfig:
 @dataclass
 class InferenceConfig:
     task: str
+    source: Union[str, int]
     nms: NMSConfig
 
 

yolo/config/task/dataset/demo.yaml

@@ -1,3 +0,0 @@
-path: demo
-
-auto_download:

yolo/config/task/inference.yaml

@@ -1,11 +1,10 @@
 task: inference
-defaults:
-  - dataset: demo
+source: demo/images/inference/image.png
 data:
   batch_size: 16
   shuffle: False
   pin_memory: True
   data_augment: {}
 nms:
-  min_confidence: 0.75
+  min_confidence: 0.1
   min_iou: 0.5

yolo/tools/data_loader.py

@@ -1,16 +1,19 @@
 import os
 from os import path
-from typing import List, Tuple, Union
+from queue import Empty, Queue
+from threading import Event, Thread
+from typing import Generator, List, Optional, Tuple, Union
 
+import cv2
 import hydra
 import numpy as np
 import torch
 from loguru import logger
 from PIL import Image
 from rich.progress import track
+from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import functional as TF
-from tqdm.rich import tqdm
 
 from yolo.config.config import Config, TrainConfig
 from yolo.tools.data_augmentation import (
@@ -199,12 +202,107 @@ class YoloDataLoader(DataLoader):
 
 
 def create_dataloader(config: Config):
+    if config.task.task == "inference":
+        return StreamDataLoader(config)
+
     if config.task.dataset.auto_download:
         prepare_dataset(config.task.dataset)
 
     return YoloDataLoader(config)
 
 
+class StreamDataLoader:
+    def __init__(self, config: Config):
+        self.source = config.task.source
+        self.running = True
+        self.is_stream = isinstance(self.source, int) or self.source.lower().startswith("rtmp://")
+
+        self.transform = AugmentationComposer([], config.image_size[0])
+        self.stop_event = Event()
+
+        if self.is_stream:
+            self.cap = cv2.VideoCapture(self.source)
+        else:
+            self.queue = Queue()
+            self.thread = Thread(target=self.load_source)
+            self.thread.start()
+
+    def load_source(self):
+        if os.path.isdir(self.source):  # image folder
+            self.load_image_folder(self.source)
+        elif any(self.source.lower().endswith(ext) for ext in [".mp4", ".avi", ".mkv"]):  # Video file
+            self.load_video_file(self.source)
+        else:  # Single image
+            self.process_image(self.source)
+
+    def load_image_folder(self, folder):
+        for root, _, files in os.walk(folder):
+            for file in files:
+                if self.stop_event.is_set():
+                    break
+                if any(file.lower().endswith(ext) for ext in [".jpg", ".jpeg", ".png", ".bmp"]):
+                    self.process_image(os.path.join(root, file))
+
+    def process_image(self, image_path):
+        image = Image.open(image_path).convert("RGB")
+        if image is None:
+            raise ValueError(f"Error loading image: {image_path}")
+        self.process_frame(image)
+
+    def load_video_file(self, video_path):
+        cap = cv2.VideoCapture(video_path)
+        while self.running:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            self.process_frame(frame)
+        cap.release()
+
+    def cv2_to_tensor(self, frame: np.ndarray) -> Tensor:
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame_float = frame_rgb.astype("float32") / 255.0
+        return torch.from_numpy(frame_float).permute(2, 0, 1)[None]
+
+    def process_frame(self, frame):
+        if isinstance(frame, np.ndarray):
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = Image.fromarray(frame)
+        frame, _ = self.transform(frame, torch.zeros(0, 5))
+        frame = TF.to_tensor(frame)[None]
+        if not self.is_stream:
+            self.queue.put(frame)
+        else:
+            self.current_frame = frame
+
+    def __iter__(self) -> Generator[Tensor, None, None]:
+        return self
+
+    def __next__(self) -> Tensor:
+        if self.is_stream:
+            ret, frame = self.cap.read()
+            if not ret:
+                self.stop()
+                raise StopIteration
+            self.process_frame(frame)
+            return self.current_frame
+        else:
+            try:
+                frame = self.queue.get(timeout=1)
+                return frame
+            except Empty:
+                raise StopIteration
+
+    def stop(self):
+        self.running = False
+        if self.is_stream:
+            self.cap.release()
+        else:
+            self.thread.join(timeout=1)
+
+    def __len__(self):
+        return self.queue.qsize() if not self.is_stream else 0
+
+
 @hydra.main(config_path="../config", config_name="config", version_base=None)
 def main(cfg):
     dataloader = create_dataloader(cfg)

yolo/tools/drawer.py

@@ -14,6 +14,7 @@ def draw_bboxes(
     *,
     scaled_bbox: bool = True,
     save_path: str = "",
+    save_name: str = "visualize.png",
 ):
     """
     Draw bounding boxes on an image.
@@ -46,7 +47,7 @@ def draw_bboxes(
         draw.rectangle(shape, outline="red", width=3)
         draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
 
-    save_image_path = os.path.join(save_path, "visualize.png")
+    save_image_path = os.path.join(save_path, save_name)
     img.save(save_image_path)  # Save the image with annotations
     logger.info(f"💾 Saved visualize image at {save_image_path}")
     return img

yolo/tools/solver.py

@@ -7,6 +7,7 @@ from torch.cuda.amp import GradScaler, autocast
 
 from yolo.config.config import Config, TrainConfig
 from yolo.model.yolo import YOLO
+from yolo.tools.data_loader import StreamDataLoader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.loss_functions import get_loss_function
 from yolo.utils.bounding_box_utils import AnchorBoxConverter, bbox_nms
@@ -103,15 +104,26 @@ class ModelTester:
         self.nms = cfg.task.nms
         self.save_path = save_path
 
-    def solve(self, dataloader):
+    def solve(self, dataloader: StreamDataLoader):
         logger.info("👀 Start Inference!")
 
-        for images, _ in dataloader:
-            images = images.to(self.device)
-            with torch.no_grad():
-                raw_output = self.model(images)
-            predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
-
-        nms_out = bbox_nms(predict, self.nms)
-        for image, bbox in zip(images, nms_out):
-            draw_bboxes(image, bbox, scaled_bbox=False, save_path=self.save_path)
+        try:
+            for idx, images in enumerate(dataloader):
+                images = images.to(self.device)
+                with torch.no_grad():
+                    raw_output = self.model(images)
+                predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
+                nms_out = bbox_nms(predict, self.nms)
+                draw_bboxes(
+                    images[0], nms_out[0], scaled_bbox=False, save_path=self.save_path, save_name=f"frame{idx:03d}.png"
+                )
+        except KeyboardInterrupt:
+            logger.error("Interrupted by user")
+            dataloader.stop_event.set()
+            dataloader.stop()
+        except Exception as e:
+            logger.error(e)
+            dataloader.stop_event.set()
+            dataloader.stop()
+            raise e
+        dataloader.stop()

yolo/utils/bounding_box_utils.py

@@ -303,7 +303,7 @@ def bbox_nms(predicts: Tensor, nms_cfg: NMSConfig):
     batch_idx, *_ = torch.where(valid_mask)
     nms_idx = batched_nms(valid_box, valid_cls, batch_idx, nms_cfg.min_iou)
     predicts_nms = []
-    for idx in range(batch_idx.max() + 1):
+    for idx in range(predicts.size(0)):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]
 
         predict_nms = torch.cat([valid_cls[instance_idx][:, None], valid_box[instance_idx]], dim=-1)