Added ability for suffix

Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
Jim Martens 2019-09-16 12:22:12 +02:00
parent 06be98f59f
commit d16edef94b
1 changed files with 4 additions and 2 deletions

View File

@ -86,7 +86,8 @@ def draw_bbox_figure(image_filename: str, labels: Sequence[np.ndarray],
instances: Sequence[Sequence[np.ndarray]],
image_size: int,
output_path: str, coco_path: str,
get_coco_cat_maps_func: callable) -> None:
get_coco_cat_maps_func: callable,
suffix: str) -> None:
"""
Draws a bounding box figure and saves it under the image file name under the output path.
@ -98,6 +99,7 @@ def draw_bbox_figure(image_filename: str, labels: Sequence[np.ndarray],
output_path: path to save the images in
coco_path: path to the COCO data set
get_coco_cat_maps_func: callable that returns the COCO category maps for a given annotation file
suffix: suffix of the saved figure file
"""
annotation_file_train = f"{coco_path}/annotations/instances_minival2014.json"
_, _, _, classes_to_names = get_coco_cat_maps_func(annotation_file_train)
@ -107,7 +109,7 @@ def draw_bbox_figure(image_filename: str, labels: Sequence[np.ndarray],
with Image.open(image_filename) as _image:
np_image = np.array(_image, dtype=np.uint8)
image = Image.fromarray(np_image)
figure_filename = f"{output_path}/{os.path.basename(image_filename)}_bboxes.png"
figure_filename = f"{output_path}/{os.path.basename(image_filename)}_bboxes_{suffix}.png"
drawables = [(colors[i], _instances) for i, _instances in enumerate(instances)]
drawables.append((colors[-9], labels))
_draw_bbox_image(image=image,