π₯ [Remove] scale function in drawer
Browse files- yolo/tools/drawer.py +11 -15
- 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("π
|
33 |
-
img = img[0]
|
34 |
-
bboxes = bboxes[0]
|
35 |
img = to_pil_image(img)
|
36 |
|
37 |
draw = ImageDraw.Draw(img, "RGBA")
|
38 |
-
|
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 |
-
|
47 |
-
|
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 |
-
|
58 |
-
|
|
|
|
|
|
|
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,
|