:bug: [Fix] a bug with drawing picture after merge

yolo/config/task/inference.yaml

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

yolo/tools/drawer.py

@@ -13,8 +13,6 @@ def draw_bboxes(
     img: Union[Image.Image, torch.Tensor],
     bboxes: List[List[Union[int, float]]],
     *,
-    save_path: str = "",
-    save_name: str = "visualize.png",
     idx2label: Optional[list],
 ):
     """
@@ -114,6 +112,6 @@ def draw_model(*, model_cfg=None, model=None, v7_base=False):
                 dot.edge(str(idx), str(jdx))
     try:
         dot.render("Model-arch", format="png", cleanup=True)
+        logger.info("🎨 Drawing Model Architecture at Model-arch.png")
     except:
-        logger.info("Warning: Could not find graphviz backend, continue without drawing the model architecture")
-    logger.info("🎨 Drawing Model Architecture at Model-arch.png")
+        logger.warning("⚠️ Could not find graphviz backend, continue without drawing the model architecture")

yolo/tools/solver.py

@@ -108,7 +108,8 @@ 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.save_path = save_path
+        self.save_predict = getattr(cfg.task, "save_predict", None)
         self.idx2label = cfg.class_list
 
     def solve(self, dataloader: StreamDataLoader):
@@ -124,27 +125,23 @@ class ModelTester:
                 images = images.to(self.device)
                 with torch.no_grad():
                     predicts = self.model(images)
-                predicts = self.vec2box(predicts["Main"])
+                    predicts = self.vec2box(predicts["Main"])
                 nms_out = bbox_nms(predicts[0], predicts[2], self.nms)
-                draw_bboxes(
-                    images[0],
-                    nms_out[0],
-                    save_path=self.save_path,
-                    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}")
+                img = draw_bboxes(images[0], nms_out[0], idx2label=self.idx2label)
 
                 if dataloader.is_stream:
-                    img = np.array(img)
-                    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-                    cv2.imshow("Result", img)
+                    img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+                    cv2.imshow("Prediction", img)
                     if cv2.waitKey(1) & 0xFF == ord("q"):
                         break
+                    if not self.save_predict:
+                        continue
+
+                if self.save_predict == False:
+                    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}")
+
         except (KeyboardInterrupt, Exception) as e:
             dataloader.stop_event.set()
             dataloader.stop()