henry000 commited on
Commit
7967aab
Β·
1 Parent(s): 4750cd0

πŸ”¨ [Update] drawer with saving path

Browse files
yolo/tools/data_augmentation.py CHANGED
@@ -38,10 +38,10 @@ class PadAndResize:
38
 
39
  resized_img = square_img.resize((self.image_size, self.image_size))
40
 
41
- boxes[:, 1] = (boxes[:, 1] + left) * scale # xmin
42
- boxes[:, 2] = (boxes[:, 2] + top) * scale # ymin
43
- boxes[:, 3] = (boxes[:, 3] + left) * scale # xmax
44
- boxes[:, 4] = (boxes[:, 4] + top) * scale # ymax
45
 
46
  return resized_img, boxes
47
 
 
38
 
39
  resized_img = square_img.resize((self.image_size, self.image_size))
40
 
41
+ boxes[:, 1] = (boxes[:, 1] + left) * scale
42
+ boxes[:, 2] = (boxes[:, 2] + top) * scale
43
+ boxes[:, 3] = (boxes[:, 3] + left) * scale
44
+ boxes[:, 4] = (boxes[:, 4] + top) * scale
45
 
46
  return resized_img, boxes
47
 
yolo/tools/drawer.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from typing import List, Union
2
 
3
  import numpy as np
@@ -8,7 +9,11 @@ from torchvision.transforms.functional import to_pil_image
8
 
9
 
10
  def draw_bboxes(
11
- img: Union[Image.Image, torch.Tensor], bboxes: List[List[Union[int, float]]], *, scaled_bbox: bool = True
 
 
 
 
12
  ):
13
  """
14
  Draw bounding boxes on an image.
@@ -21,7 +26,7 @@ def draw_bboxes(
21
  # Convert tensor image to PIL Image if necessary
22
  if isinstance(img, torch.Tensor):
23
  if img.dim() > 3:
24
- logger.info("Multi-frame tensor detected, using the first image.")
25
  img = img[0]
26
  bboxes = bboxes[0]
27
  img = to_pil_image(img)
@@ -41,8 +46,9 @@ def draw_bboxes(
41
  draw.rectangle(shape, outline="red", width=3)
42
  draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
43
 
44
- img.save("visualize.jpg") # Save the image with annotations
45
- logger.info("Saved visualize image at visualize.png")
 
46
  return img
47
 
48
 
 
1
+ import os
2
  from typing import List, Union
3
 
4
  import numpy as np
 
9
 
10
 
11
  def draw_bboxes(
12
+ img: Union[Image.Image, torch.Tensor],
13
+ bboxes: List[List[Union[int, float]]],
14
+ *,
15
+ scaled_bbox: bool = True,
16
+ save_path: str = "",
17
  ):
18
  """
19
  Draw bounding boxes on an image.
 
26
  # Convert tensor image to PIL Image if necessary
27
  if isinstance(img, torch.Tensor):
28
  if img.dim() > 3:
29
+ logger.warning("πŸ” Multi-frame tensor detected, using the first image.")
30
  img = img[0]
31
  bboxes = bboxes[0]
32
  img = to_pil_image(img)
 
46
  draw.rectangle(shape, outline="red", width=3)
47
  draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
48
 
49
+ save_image_path = os.path.join(save_path, "visualize.png")
50
+ img.save(save_image_path) # Save the image with annotations
51
+ logger.info(f"πŸ’Ύ Saved visualize image at {save_image_path}")
52
  return img
53
 
54