Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 1 | # SPDX-FileCopyrightText: Copyright 2021-2023 Arm Limited and/or its affiliates <open-source-office@arm.com> |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. |
| 6 | # You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. |
| 15 | |
| 16 | """ |
| 17 | Utility script to convert a set of RGB images in a given location into |
| 18 | corresponding cpp files and a single hpp file referencing the vectors |
| 19 | from the cpp files. |
| 20 | """ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 21 | import glob |
| 22 | import math |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 23 | import typing |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 24 | from argparse import ArgumentParser |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 25 | from dataclasses import dataclass |
| 26 | from pathlib import Path |
Richard Burton | 1706962 | 2022-03-17 10:54:26 +0000 | [diff] [blame] | 27 | |
| 28 | import numpy as np |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 29 | from PIL import Image, UnidentifiedImageError |
| 30 | from jinja2 import Environment, FileSystemLoader |
| 31 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 32 | from gen_utils import GenUtils |
| 33 | |
| 34 | # pylint: disable=duplicate-code |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 35 | parser = ArgumentParser() |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 36 | |
| 37 | parser.add_argument( |
| 38 | "--image_path", |
| 39 | type=str, |
| 40 | help="path to images folder or image file to convert." |
| 41 | ) |
| 42 | |
| 43 | parser.add_argument( |
| 44 | "--source_folder_path", |
| 45 | type=str, |
| 46 | help="path to source folder to be generated." |
| 47 | ) |
| 48 | |
| 49 | parser.add_argument( |
| 50 | "--header_folder_path", |
| 51 | type=str, |
| 52 | help="path to header folder to be generated." |
| 53 | ) |
| 54 | |
| 55 | parser.add_argument( |
| 56 | "--image_size", |
| 57 | type=int, |
| 58 | nargs=2, |
| 59 | help="Size (width and height) of the converted images." |
| 60 | ) |
| 61 | |
| 62 | parser.add_argument( |
| 63 | "--license_template", |
| 64 | type=str, |
| 65 | help="Header template file", |
| 66 | default="header_template.txt" |
| 67 | ) |
| 68 | |
| 69 | parsed_args = parser.parse_args() |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 70 | |
Richard Burton | 1706962 | 2022-03-17 10:54:26 +0000 | [diff] [blame] | 71 | env = Environment(loader=FileSystemLoader(Path(__file__).parent / 'templates'), |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 72 | trim_blocks=True, |
| 73 | lstrip_blocks=True) |
| 74 | |
| 75 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 76 | # pylint: enable=duplicate-code |
| 77 | @dataclass |
| 78 | class ImagesParams: |
| 79 | """ |
| 80 | Template params for Images.hpp and Images.cc |
| 81 | """ |
| 82 | num_images: int |
| 83 | image_size: typing.Sequence |
| 84 | image_array_names: typing.List[str] |
| 85 | image_filenames: typing.List[str] |
| 86 | |
| 87 | |
| 88 | def write_hpp_file( |
| 89 | images_params: ImagesParams, |
| 90 | header_file_path: Path, |
| 91 | cc_file_path: Path, |
| 92 | header_template_file: str, |
| 93 | ): |
| 94 | """ |
| 95 | Write Images.hpp and Images.cc |
| 96 | |
| 97 | @param images_params: Template params |
| 98 | @param header_file_path: Images.hpp path |
| 99 | @param cc_file_path: Images.cc path |
| 100 | @param header_template_file: Header template file name |
| 101 | """ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 102 | print(f"++ Generating {header_file_path}") |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 103 | hdr = GenUtils.gen_header(env, header_template_file) |
| 104 | |
| 105 | image_size = str(images_params.image_size[0] * images_params.image_size[1] * 3) |
| 106 | |
| 107 | env \ |
| 108 | .get_template('Images.hpp.template') \ |
| 109 | .stream(common_template_header=hdr, |
| 110 | imgs_count=images_params.num_images, |
| 111 | img_size=image_size, |
| 112 | var_names=images_params.image_array_names) \ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 113 | .dump(str(header_file_path)) |
| 114 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 115 | env \ |
| 116 | .get_template('Images.cc.template') \ |
| 117 | .stream(common_template_header=hdr, |
| 118 | var_names=images_params.image_array_names, |
| 119 | img_names=images_params.image_filenames) \ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 120 | .dump(str(cc_file_path)) |
| 121 | |
| 122 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 123 | def resize_crop_image( |
| 124 | original_image: Image.Image, |
| 125 | image_size: typing.Sequence |
| 126 | ) -> np.ndarray: |
| 127 | """ |
| 128 | Resize and crop input image |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 129 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 130 | @param original_image: Image to resize and crop |
| 131 | @param image_size: New image size |
| 132 | @return: Resized and cropped image |
| 133 | """ |
Isabella Gottardi | 79d4154 | 2021-10-20 15:52:32 +0100 | [diff] [blame] | 134 | # IFM size |
| 135 | ifm_width = image_size[0] |
| 136 | ifm_height = image_size[1] |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 137 | |
Isabella Gottardi | 79d4154 | 2021-10-20 15:52:32 +0100 | [diff] [blame] | 138 | # Aspect ratio resize |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 139 | scale_ratio = (float(max(ifm_width, ifm_height)) |
| 140 | / float(min(original_image.size[0], original_image.size[1]))) |
| 141 | resized_width = int(original_image.size[0] * scale_ratio) |
| 142 | resized_height = int(original_image.size[1] * scale_ratio) |
| 143 | resized_image = original_image.resize( |
| 144 | size=(resized_width, resized_height), |
| 145 | resample=Image.Resampling.BILINEAR |
| 146 | ) |
Isabella Gottardi | 79d4154 | 2021-10-20 15:52:32 +0100 | [diff] [blame] | 147 | |
| 148 | # Crop the center of the image |
| 149 | resized_image = resized_image.crop(( |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 150 | (resized_width - ifm_width) / 2, # left |
| 151 | (resized_height - ifm_height) / 2, # top |
| 152 | (resized_width + ifm_width) / 2, # right |
Isabella Gottardi | 79d4154 | 2021-10-20 15:52:32 +0100 | [diff] [blame] | 153 | (resized_height + ifm_height) / 2 # bottom |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 154 | )) |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 155 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 156 | return np.array(resized_image, dtype=np.uint8).flatten() |
| 157 | |
| 158 | |
| 159 | def write_individual_img_cc_file( |
| 160 | rgb_data: np.ndarray, |
| 161 | image_filename: str, |
| 162 | cc_filename: Path, |
| 163 | header_template_file: str, |
| 164 | array_name: str |
| 165 | ): |
| 166 | """ |
| 167 | Write image.cc |
| 168 | |
| 169 | @param rgb_data: Image data |
| 170 | @param image_filename: Image file name |
| 171 | @param cc_filename: image.cc path |
| 172 | @param header_template_file: Header template file name |
| 173 | @param array_name: C++ array name |
| 174 | """ |
| 175 | print(f"++ Converting {image_filename} to {cc_filename.name}") |
| 176 | |
| 177 | hdr = GenUtils.gen_header(env, header_template_file, image_filename) |
| 178 | |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 179 | hex_line_generator = (', '.join(map(hex, sub_arr)) |
| 180 | for sub_arr in np.array_split(rgb_data, math.ceil(len(rgb_data) / 20))) |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 181 | env \ |
| 182 | .get_template('image.cc.template') \ |
| 183 | .stream(common_template_header=hdr, |
| 184 | var_name=array_name, |
| 185 | img_data=hex_line_generator) \ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 186 | .dump(str(cc_filename)) |
| 187 | |
| 188 | |
| 189 | def main(args): |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 190 | """ |
| 191 | Convert images |
| 192 | @param args: Parsed args |
| 193 | """ |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 194 | # Keep the count of the images converted |
| 195 | image_idx = 0 |
| 196 | image_filenames = [] |
| 197 | image_array_names = [] |
| 198 | |
Richard Burton | 1706962 | 2022-03-17 10:54:26 +0000 | [diff] [blame] | 199 | if Path(args.image_path).is_dir(): |
| 200 | filepaths = sorted(glob.glob(str(Path(args.image_path) / '**/*.*'), recursive=True)) |
| 201 | elif Path(args.image_path).is_file(): |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 202 | filepaths = [args.image_path] |
| 203 | else: |
| 204 | raise OSError("Directory or file does not exist.") |
| 205 | |
| 206 | for filepath in filepaths: |
Richard Burton | 1706962 | 2022-03-17 10:54:26 +0000 | [diff] [blame] | 207 | filename = Path(filepath).name |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 208 | |
| 209 | try: |
| 210 | original_image = Image.open(filepath).convert("RGB") |
| 211 | except UnidentifiedImageError: |
| 212 | print(f"-- Skipping file {filepath} due to unsupported image format.") |
| 213 | continue |
| 214 | |
| 215 | image_filenames.append(filename) |
| 216 | |
| 217 | # Save the cc file |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 218 | cc_filename = (Path(args.source_folder_path) / |
| 219 | (Path(filename).stem.replace(" ", "_") + ".cc")) |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 220 | array_name = "im" + str(image_idx) |
| 221 | image_array_names.append(array_name) |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 222 | |
| 223 | rgb_data = resize_crop_image(original_image, args.image_size) |
| 224 | write_individual_img_cc_file( |
| 225 | rgb_data, filename, cc_filename, args.license_template, array_name |
| 226 | ) |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 227 | |
| 228 | # Increment image index |
| 229 | image_idx = image_idx + 1 |
| 230 | |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 231 | header_filepath = Path(args.header_folder_path) / "InputFiles.hpp" |
| 232 | common_cc_filepath = Path(args.source_folder_path) / "InputFiles.cc" |
| 233 | |
| 234 | images_params = ImagesParams(image_idx, args.image_size, image_array_names, image_filenames) |
Kshitij Sisodia | 659fcd9 | 2021-05-19 10:30:06 +0100 | [diff] [blame] | 235 | |
| 236 | if len(image_filenames) > 0: |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 237 | write_hpp_file(images_params, header_filepath, common_cc_filepath, args.license_template) |
Kshitij Sisodia | 659fcd9 | 2021-05-19 10:30:06 +0100 | [diff] [blame] | 238 | else: |
| 239 | raise FileNotFoundError("No valid images found.") |
alexander | 3c79893 | 2021-03-26 21:42:19 +0000 | [diff] [blame] | 240 | |
| 241 | |
| 242 | if __name__ == '__main__': |
Alex Tawse | daba3cf | 2023-09-29 15:55:38 +0100 | [diff] [blame] | 243 | main(parsed_args) |