Creates batches of pre-processed images.
| 24 | |
| 25 | |
| 26 | class ImageBatcher: |
| 27 | """ |
| 28 | Creates batches of pre-processed images. |
| 29 | """ |
| 30 | |
| 31 | def __init__(self, input, shape, dtype, max_num_images=None, exact_batches=False, preprocessor="EfficientDet", shuffle_files=False): |
| 32 | """ |
| 33 | :param input: The input directory to read images from. |
| 34 | :param shape: The tensor shape of the batch to prepare, either in NCHW or NHWC format. |
| 35 | :param dtype: The (numpy) datatype to cast the batched data to. |
| 36 | :param max_num_images: The maximum number of images to read from the directory. |
| 37 | :param exact_batches: This defines how to handle a number of images that is not an exact multiple of the batch |
| 38 | size. If false, it will pad the final batch with zeros to reach the batch size. If true, it will *remove* the |
| 39 | last few images in excess of a batch size multiple, to guarantee batches are exact (useful for calibration). |
| 40 | :param preprocessor: Set the preprocessor to use, depending on which network is being used. |
| 41 | :param shuffle_files: Shuffle the list of files before batching. |
| 42 | """ |
| 43 | # Find images in the given input path |
| 44 | input = os.path.realpath(input) |
| 45 | self.images = [] |
| 46 | |
| 47 | extensions = [".jpg", ".jpeg", ".png", ".bmp"] |
| 48 | |
| 49 | def is_image(path): |
| 50 | return os.path.isfile(path) and os.path.splitext(path)[1].lower() in extensions |
| 51 | |
| 52 | if os.path.isdir(input): |
| 53 | self.images = [os.path.join(input, f) for f in os.listdir(input) if is_image(os.path.join(input, f))] |
| 54 | self.images.sort() |
| 55 | if shuffle_files: |
| 56 | random.seed(47) |
| 57 | random.shuffle(self.images) |
| 58 | elif os.path.isfile(input): |
| 59 | if is_image(input): |
| 60 | self.images.append(input) |
| 61 | self.num_images = len(self.images) |
| 62 | if self.num_images < 1: |
| 63 | print("No valid {} images found in {}".format("/".join(extensions), input)) |
| 64 | sys.exit(1) |
| 65 | |
| 66 | # Handle Tensor Shape |
| 67 | self.dtype = dtype |
| 68 | self.shape = shape |
| 69 | assert len(self.shape) == 4 |
| 70 | self.batch_size = shape[0] |
| 71 | assert self.batch_size > 0 |
| 72 | self.format = None |
| 73 | self.width = -1 |
| 74 | self.height = -1 |
| 75 | if self.shape[1] == 3: |
| 76 | self.format = "NCHW" |
| 77 | self.height = self.shape[2] |
| 78 | self.width = self.shape[3] |
| 79 | elif self.shape[3] == 3: |
| 80 | self.format = "NHWC" |
| 81 | self.height = self.shape[1] |
| 82 | self.width = self.shape[2] |
| 83 | assert all([self.format, self.width > 0, self.height > 0]) |
no outgoing calls
no test coverage detected