henry000 commited on
Commit
3b24b35
Β·
1 Parent(s): 0d25eaa

πŸ”₯ [Remove] scale function in drawer

Browse files
Files changed (2) hide show
  1. yolo/tools/drawer.py +11 -15
  2. yolo/tools/solver.py +0 -1
yolo/tools/drawer.py CHANGED
@@ -13,7 +13,6 @@ def draw_bboxes(
13
  img: Union[Image.Image, torch.Tensor],
14
  bboxes: List[List[Union[int, float]]],
15
  *,
16
- scaled_bbox: bool = True,
17
  save_path: str = "",
18
  save_name: str = "visualize.png",
19
  idx2label: Optional[list],
@@ -29,13 +28,12 @@ def draw_bboxes(
29
  # Convert tensor image to PIL Image if necessary
30
  if isinstance(img, torch.Tensor):
31
  if img.dim() > 3:
32
- logger.warning("πŸ” Multi-frame tensor detected, using the first image.")
33
- img = img[0]
34
- bboxes = bboxes[0]
35
  img = to_pil_image(img)
36
 
37
  draw = ImageDraw.Draw(img, "RGBA")
38
- width, height = img.size
39
  try:
40
  font = ImageFont.truetype("arial.ttf", 15)
41
  except IOError:
@@ -43,19 +41,16 @@ def draw_bboxes(
43
 
44
  for bbox in bboxes:
45
  class_id, x_min, y_min, x_max, y_max, *conf = [float(val) for val in bbox]
46
- if scaled_bbox:
47
- x_min = x_min * width
48
- x_max = x_max * width
49
- y_min = y_min * height
50
- y_max = y_max * height
51
- shape = [(x_min, y_min), (x_max, y_max)]
52
  random.seed(int(class_id))
53
  color_map = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
54
- draw.rounded_rectangle(shape, outline=(*color_map, 170), radius=5)
55
- draw.rounded_rectangle(shape, fill=(*color_map, 50), radius=5)
56
 
57
- text_class = str(idx2label[int(class_id)] if idx2label else class_id)
58
- label_text = f"{text_class}" + (f" {conf[0]: .0%}" if conf else "")
 
 
 
59
 
60
  text_bbox = font.getbbox(label_text)
61
  text_width = text_bbox[2] - text_bbox[0]
@@ -65,6 +60,7 @@ def draw_bboxes(
65
  draw.rounded_rectangle(text_background, fill=(*color_map, 175), radius=2)
66
  draw.text((x_min, y_min), label_text, fill="white", font=font)
67
 
 
68
  save_image_path = os.path.join(save_path, save_name)
69
  img.save(save_image_path) # Save the image with annotations
70
  logger.info(f"πŸ’Ύ Saved visualize image at {save_image_path}")
 
13
  img: Union[Image.Image, torch.Tensor],
14
  bboxes: List[List[Union[int, float]]],
15
  *,
 
16
  save_path: str = "",
17
  save_name: str = "visualize.png",
18
  idx2label: Optional[list],
 
28
  # Convert tensor image to PIL Image if necessary
29
  if isinstance(img, torch.Tensor):
30
  if img.dim() > 3:
31
+ logger.warning("πŸ” >3 dimension tensor detected, using the 0-idx image.")
32
+ img, bboxes = img[0], bboxes[0]
 
33
  img = to_pil_image(img)
34
 
35
  draw = ImageDraw.Draw(img, "RGBA")
36
+
37
  try:
38
  font = ImageFont.truetype("arial.ttf", 15)
39
  except IOError:
 
41
 
42
  for bbox in bboxes:
43
  class_id, x_min, y_min, x_max, y_max, *conf = [float(val) for val in bbox]
44
+ bbox = [(x_min, y_min), (x_max, y_max)]
45
+
 
 
 
 
46
  random.seed(int(class_id))
47
  color_map = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
 
 
48
 
49
+ draw.rounded_rectangle(bbox, outline=(*color_map, 170), radius=5)
50
+ draw.rounded_rectangle(bbox, fill=(*color_map, 50), radius=5)
51
+
52
+ class_text = str(idx2label[int(class_id)] if idx2label else class_id)
53
+ label_text = f"{class_text}" + (f" {conf[0]: .0%}" if conf else "")
54
 
55
  text_bbox = font.getbbox(label_text)
56
  text_width = text_bbox[2] - text_bbox[0]
 
60
  draw.rounded_rectangle(text_background, fill=(*color_map, 175), radius=2)
61
  draw.text((x_min, y_min), label_text, fill="white", font=font)
62
 
63
+ os.makedirs(save_path, exist_ok=True)
64
  save_image_path = os.path.join(save_path, save_name)
65
  img.save(save_image_path) # Save the image with annotations
66
  logger.info(f"πŸ’Ύ Saved visualize image at {save_image_path}")
yolo/tools/solver.py CHANGED
@@ -117,7 +117,6 @@ class ModelTester:
117
  draw_bboxes(
118
  images[0],
119
  nms_out[0],
120
- scaled_bbox=False,
121
  save_path=self.save_path,
122
  save_name=f"frame{idx:03d}.png",
123
  idx2label=self.idx2label,
 
117
  draw_bboxes(
118
  images[0],
119
  nms_out[0],
 
120
  save_path=self.save_path,
121
  save_name=f"frame{idx:03d}.png",
122
  idx2label=self.idx2label,