🔨 [Update] drawer with saving path

yolo/tools/data_augmentation.py

@@ -38,10 +38,10 @@ class PadAndResize:
 
         resized_img = square_img.resize((self.image_size, self.image_size))
 
-        boxes[:, 1] = (boxes[:, 1] + left) * scale  # xmin
-        boxes[:, 2] = (boxes[:, 2] + top) * scale  # ymin
-        boxes[:, 3] = (boxes[:, 3] + left) * scale  # xmax
-        boxes[:, 4] = (boxes[:, 4] + top) * scale  # ymax
+        boxes[:, 1] = (boxes[:, 1] + left) * scale
+        boxes[:, 2] = (boxes[:, 2] + top) * scale
+        boxes[:, 3] = (boxes[:, 3] + left) * scale
+        boxes[:, 4] = (boxes[:, 4] + top) * scale
 
         return resized_img, boxes
 

yolo/tools/drawer.py

@@ -1,3 +1,4 @@
+import os
 from typing import List, Union
 
 import numpy as np
@@ -8,7 +9,11 @@ from torchvision.transforms.functional import to_pil_image
 
 
 def draw_bboxes(
-    img: Union[Image.Image, torch.Tensor], bboxes: List[List[Union[int, float]]], *, scaled_bbox: bool = True
+    img: Union[Image.Image, torch.Tensor],
+    bboxes: List[List[Union[int, float]]],
+    *,
+    scaled_bbox: bool = True,
+    save_path: str = "",
 ):
     """
     Draw bounding boxes on an image.
@@ -21,7 +26,7 @@ def draw_bboxes(
     # Convert tensor image to PIL Image if necessary
     if isinstance(img, torch.Tensor):
         if img.dim() > 3:
-            logger.info("Multi-frame tensor detected, using the first image.")
+            logger.warning("🔍 Multi-frame tensor detected, using the first image.")
             img = img[0]
             bboxes = bboxes[0]
         img = to_pil_image(img)
@@ -41,8 +46,9 @@ def draw_bboxes(
         draw.rectangle(shape, outline="red", width=3)
         draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
 
-    img.save("visualize.jpg")  # Save the image with annotations
-    logger.info("Saved visualize image at visualize.png")
+    save_image_path = os.path.join(save_path, "visualize.png")
+    img.save(save_image_path)  # Save the image with annotations
+    logger.info(f"💾 Saved visualize image at {save_image_path}")
     return img