File size: 4,836 Bytes
00c3521 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
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):
# parse raw img
raw_image = cv2.imread(str(png_path), cv2.IMREAD_UNCHANGED)
# parse metadata
metadata = json_read(png_path.with_suffix('.json'),
object_hook=fraction_from_json)
start_time = time.perf_counter()
# executing img pipelinex
pipeline_exec = PipelineExecutor(raw_image,
metadata,
os.path.basename(
str(png_path)).split('.')[0],
self.pipeline_demo,
save_dir=self.save_dir)
# process img
output_image = pipeline_exec()
end_time = time.perf_counter()
self.process_times.append(end_time - start_time)
# save results
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))
|