[🎨] Add support for displaying webcam videos with predicted bounding boxes (#27)

* ✅ [Pass] tests, skip drawing if graphviz not found

* ✨ [Add] Display processed webcam videos

---------

Co-authored-by: HenryTsui <[email protected]>

yolo/config/config.py

@@ -108,6 +108,7 @@ class InferenceConfig:
     nms: NMSConfig
     data: DataConfig
     fast_inference: Optional[None]
+    save_predict: bool
 
 
 @dataclass

yolo/config/task/inference.yaml

@@ -7,4 +7,5 @@ data:
   data_augment: {}
 nms:
   min_confidence: 0.5
-  min_iou: 0.5
+  min_iou: 0.5
+save_predict: true

yolo/tools/drawer.py

@@ -2,6 +2,7 @@ import os
 import random
 from typing import List, Optional, Union
 
+
 import numpy as np
 import torch
 from loguru import logger
@@ -65,9 +66,6 @@ def draw_bboxes(
         draw.rounded_rectangle(text_background, fill=(*color_map, 175), radius=2)
         draw.text((x_min, y_min), label_text, fill="white", font=font)
 
-    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

@@ -1,3 +1,5 @@
+import os
+
 import torch
 from loguru import logger
 from torch import Tensor
@@ -106,12 +108,15 @@ class ModelTester:
 
         self.anchor2box = AnchorBoxConverter(cfg.model, cfg.image_size, device)
         self.nms = cfg.task.nms
+        self.save_path = save_path if getattr(cfg.task, "save_predict", True) else None
         self.idx2label = cfg.class_list
-        self.save_path = save_path
 
     def solve(self, dataloader: StreamDataLoader):
         logger.info("👀 Start Inference!")
 
+        if dataloader.is_stream:
+            import cv2
+            import numpy as np
         try:
             for idx, images in enumerate(dataloader):
                 images = images.to(self.device)
@@ -119,7 +124,7 @@ class ModelTester:
                     raw_output = self.model(images)
                 predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
                 nms_out = bbox_nms(predict, self.nms)
-                draw_bboxes(
+                img = draw_bboxes(
                     images[0],
                     nms_out[0],
                     scaled_bbox=False,
@@ -127,6 +132,18 @@ class ModelTester:
                     save_name=f"frame{idx:03d}.png",
                     idx2label=self.idx2label,
                 )
+                logger.info(f"img size: {img.shape}")
+                if self.save_path is not None:
+                    save_image_path = os.path.join(self.save_path, f"frame{idx:03d}.png")
+                    img.save(save_image_path)
+                    logger.info(f"💾 Saved visualize image at {save_image_path}")
+
+                if dataloader.is_stream:
+                    img = np.array(img)
+                    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+                    cv2.imshow("Result", img)
+                    if cv2.waitKey(1) & 0xFF == ord("q"):
+                        break
         except (KeyboardInterrupt, Exception) as e:
             dataloader.stop_event.set()
             dataloader.stop()