Creates batches of pre-processed images.
| 28 | sys.exit(1) |
| 29 | |
| 30 | class ImageBatcher: |
| 31 | """ |
| 32 | Creates batches of pre-processed images. |
| 33 | """ |
| 34 | |
| 35 | def __init__(self, input, shape, dtype, max_num_images=None, exact_batches=False, config_file=None): |
| 36 | """ |
| 37 | :param input: The input directory to read images from. |
| 38 | :param shape: The tensor shape of the batch to prepare, either in NCHW or NHWC format. |
| 39 | :param dtype: The (numpy) datatype to cast the batched data to. |
| 40 | :param max_num_images: The maximum number of images to read from the directory. |
| 41 | :param exact_batches: This defines how to handle a number of images that is not an exact multiple of the batch |
| 42 | size. If false, it will pad the final batch with zeros to reach the batch size. If true, it will *remove* the |
| 43 | last few images in excess of a batch size multiple, to guarantee batches are exact (useful for calibration). |
| 44 | :param config_file: The path pointing to the Detectron 2 yaml file which describes the model. |
| 45 | """ |
| 46 | |
| 47 | def det2_setup(config_file): |
| 48 | """ |
| 49 | Create configs and perform basic setups. |
| 50 | """ |
| 51 | cfg = get_cfg() |
| 52 | if config_file is not None: |
| 53 | cfg.merge_from_file(config_file) |
| 54 | cfg.freeze() |
| 55 | return cfg |
| 56 | |
| 57 | # Set up Detectron 2 model configuration. |
| 58 | self.det2_cfg = det2_setup(config_file) |
| 59 | |
| 60 | # Extract min and max dimensions for testing. |
| 61 | self.min_size_test = self.det2_cfg.INPUT.MIN_SIZE_TEST |
| 62 | self.max_size_test = self.det2_cfg.INPUT.MAX_SIZE_TEST |
| 63 | |
| 64 | # Find images in the given input path. |
| 65 | input = os.path.realpath(input) |
| 66 | self.images = [] |
| 67 | |
| 68 | extensions = [".jpg", ".jpeg", ".png", ".bmp", ".ppm"] |
| 69 | |
| 70 | def is_image(path): |
| 71 | return os.path.isfile(path) and os.path.splitext(path)[1].lower() in extensions |
| 72 | |
| 73 | if os.path.isdir(input): |
| 74 | self.images = [os.path.join(input, f) for f in os.listdir(input) if is_image(os.path.join(input, f))] |
| 75 | self.images.sort() |
| 76 | elif os.path.isfile(input): |
| 77 | if is_image(input): |
| 78 | self.images.append(input) |
| 79 | self.num_images = len(self.images) |
| 80 | if self.num_images < 1: |
| 81 | print("No valid {} images found in {}".format("/".join(extensions), input)) |
| 82 | sys.exit(1) |
| 83 | |
| 84 | # Handle Tensor Shape. |
| 85 | self.dtype = dtype |
| 86 | self.shape = shape |
| 87 | assert len(self.shape) == 4 |
no outgoing calls
no test coverage detected