🔀 [Merge] branch 'INFERENCE' of github.com:WongKinYiu/yolov9mit into INFERENCE

yolo/config/config.py

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

yolo/config/task/inference.yaml

@@ -8,3 +8,4 @@ data:
 nms:
   min_confidence: 0.5
   min_iou: 0.5
+save_predict: true

yolo/tools/drawer.py

@@ -60,10 +60,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)
 
-    os.makedirs(save_path, exist_ok=True)
-    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,13 +108,17 @@ class ModelTester:
         self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
 
         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 isinstance(self.model, torch.nn.Module):
             self.model.eval()
+
+        if dataloader.is_stream:
+            import cv2
+            import numpy as np
         try:
             for idx, images in enumerate(dataloader):
                 images = images.to(self.device)
@@ -127,6 +133,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()