From 2f03be42dd0855d17ada9dc0a09424ede573b855 Mon Sep 17 00:00:00 2001 From: Jim Martens Date: Mon, 15 Jul 2019 13:25:47 +0200 Subject: [PATCH] Changed save_ssd_train_images func to work with filenames as well Signed-off-by: Jim Martens --- src/twomartens/masterthesis/debug.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/twomartens/masterthesis/debug.py b/src/twomartens/masterthesis/debug.py index 5b6bbfc..bbbbd3d 100644 --- a/src/twomartens/masterthesis/debug.py +++ b/src/twomartens/masterthesis/debug.py @@ -23,6 +23,7 @@ Functions: saves the first batch of SSD train images with overlaid ground truth bounding boxes """ import os +from typing import Union, Sequence import math import numpy as np @@ -30,7 +31,7 @@ from matplotlib import pyplot from PIL import Image -def save_ssd_train_images(images: np.ndarray, labels: np.ndarray, +def save_ssd_train_images(images: Union[np.ndarray, Sequence[str]], labels: np.ndarray, output_path: str, coco_path: str, image_size: int, get_coco_cat_maps_func: callable, custom_string: str = None) -> None: @@ -40,7 +41,7 @@ def save_ssd_train_images(images: np.ndarray, labels: np.ndarray, The images are saved both in a raw version and with bounding boxes printed on them. Args: - images: a NumPy array of images + images: a NumPy array of images or a list of filenames labels: a NumPy array of labels output_path: path to save the images in coco_path: path to the COCO data set @@ -60,7 +61,11 @@ def save_ssd_train_images(images: np.ndarray, labels: np.ndarray, for i, train_image in enumerate(images): instances = labels[i] - image = Image.fromarray(train_image) + if type(train_image) is str: + with Image.open(train_image) as _image: + image = _image + else: + image = Image.fromarray(train_image) image.save(f"{output_path}/" f"{custom_string}train_image{str(i).zfill(nr_digits)}.png")