| import argparse |
| import os |
| import time |
| import warnings |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| from tqdm import tqdm |
|
|
| from raw_prc_pipeline import (expected_img_ext, expected_landscape_img_height, |
| expected_landscape_img_width) |
| from raw_prc_pipeline.pipeline import (PipelineExecutor, |
| RawProcessingPipelineDemo) |
| from utils import fraction_from_json, json_read |
|
|
| warnings.filterwarnings("ignore") |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser( |
| description= |
| 'Demo script for processing PNG images with given metadata files.') |
| parser.add_argument( |
| '-p', |
| '--png_dir', |
| type=str, |
| default='data', |
| help='Path of the directory containing PNG images with metadata files') |
| parser.add_argument( |
| '-o', |
| '--out_dir', |
| type=str, |
| default=None, |
| help= |
| 'Path to the directory where processed images will be saved. Images will be saved in JPG format.' |
| ) |
| parser.add_argument( |
| '-ie', |
| '--illumination_estimation', |
| type=str, |
| default='gw', |
| help= |
| 'Options for illumination estimation algorithms: "gw", "wp", "sog", "iwp".' |
| ) |
| parser.add_argument( |
| '-tm', |
| '--tone_mapping', |
| type=str, |
| default='Storm', |
| help= |
| 'Options for tone mapping algorithms: "Base", "Flash", "Storm", "Linear", "Drago", "Mantiuk", "Reinhard".' |
| ) |
| parser.add_argument( |
| '-n', |
| '--denoising_flg', |
| action='store_false', |
| help= |
| 'Denoising flag. By default resulted images will be denoised with some default parameters.' |
| ) |
| parser.add_argument('-m', |
| '--camera_matrix', |
| type=float, |
| nargs=9, |
| default=[ |
| 1.06835938, -0.29882812, -0.14257812, -0.43164062, |
| 1.35546875, 0.05078125, -0.1015625, 0.24414062, |
| 0.5859375 |
| ], |
| help='Mean color matrix of Hauwei Mate 40 Pro') |
| args = parser.parse_args() |
|
|
| if args.out_dir is None: |
| args.out_dir = args.png_dir |
|
|
| return args |
|
|
|
|
| class PNGProcessingDemo: |
| def __init__(self, ie_method, tone_mapping, denoising_flg, camera_matrix, |
| save_dir): |
| self.camera_matrix = camera_matrix |
| self.save_dir = save_dir |
| self.pipeline_demo = RawProcessingPipelineDemo( |
| illumination_estimation=ie_method, |
| denoise_flg=denoising_flg, |
| tone_mapping=tone_mapping, |
| out_landscape_height=expected_landscape_img_height, |
| out_landscape_width=expected_landscape_img_width) |
| self.process_times = [] |
|
|
| def __call__(self, png_path: Path, out_path: Path): |
| |
| raw_image = cv2.imread(str(png_path), cv2.IMREAD_UNCHANGED) |
| |
| metadata = json_read(png_path.with_suffix('.json'), |
| object_hook=fraction_from_json) |
| start_time = time.perf_counter() |
| |
| pipeline_exec = PipelineExecutor(raw_image, |
| metadata, |
| os.path.basename( |
| str(png_path)).split('.')[0], |
| self.pipeline_demo, |
| save_dir=self.save_dir) |
| |
| output_image = pipeline_exec() |
| end_time = time.perf_counter() |
| self.process_times.append(end_time - start_time) |
|
|
| |
| output_image = cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR) |
| cv2.imwrite(str(out_path), output_image, |
| [cv2.IMWRITE_JPEG_QUALITY, 100]) |
|
|
|
|
| def main(png_dir, out_dir, illumination_estimation, tone_mapping, |
| denoising_flg, camera_matrix): |
| png_dir = Path(png_dir) |
| out_dir = Path(out_dir) |
| out_dir.mkdir(exist_ok=True) |
|
|
| png_paths = list(png_dir.glob('*.png')) |
| out_paths = [ |
| out_dir / png_path.with_suffix(expected_img_ext).name |
| for png_path in png_paths |
| ] |
|
|
| png_processor = PNGProcessingDemo(illumination_estimation, tone_mapping, |
| denoising_flg, camera_matrix, |
| str(out_dir)) |
|
|
| for png_path, out_path in tqdm(zip(png_paths, out_paths), |
| total=len(png_paths)): |
| png_processor(png_path, out_path) |
| print("Average processing time: {:.2f}s".format( |
| np.mean(png_processor.process_times))) |
|
|
|
|
| if __name__ == '__main__': |
| args = parse_args() |
| main(**vars(args)) |
|
|
