π¨ [Update] drawer with saving path
Browse files- yolo/tools/data_augmentation.py +4 -4
- yolo/tools/drawer.py +10 -4
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
|
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 |
|
|
|
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],
|
|
|
|
|
|
|
|
|
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.
|
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 |
-
|
45 |
-
|
|
|
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 |
|