-
Notifications
You must be signed in to change notification settings - Fork 67
/
demo.py
80 lines (56 loc) · 2.14 KB
/
demo.py
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
import os
import os.path as op
import cv2
import numpy as np
import skimage.io
from pipeline import Pipeline
from utils.yacs import Config
OUTPUT_DIR = './output'
os.makedirs(OUTPUT_DIR, exist_ok=True)
def demo_test_raw():
cfg = Config('configs/test.yaml')
pipeline = Pipeline(cfg)
raw_path = 'raw/test.RAW'
bayer = np.fromfile(raw_path, dtype='uint16', sep='')
bayer = bayer.reshape((cfg.hardware.raw_height, cfg.hardware.raw_width))
data, _ = pipeline.execute(bayer)
output_path = op.join(OUTPUT_DIR, 'test.png')
output = cv2.cvtColor(data['output'], cv2.COLOR_RGB2BGR)
cv2.imwrite(output_path, output)
def demo_nikon_d3x():
cfg = Config('configs/nikon_d3x.yaml')
pipeline = Pipeline(cfg)
pgm_path = 'raw/color_checker.pgm'
bayer = skimage.io.imread(pgm_path).astype(np.uint16)
data, _ = pipeline.execute(bayer)
output_path = op.join(OUTPUT_DIR, 'color_checker.png')
output = cv2.cvtColor(data['output'], cv2.COLOR_RGB2BGR)
cv2.imwrite(output_path, output)
def demo_intermediate_results():
cfg = Config('configs/nikon_d3x.yaml')
# Boost parameters for exaggeration effect
with cfg.unfreeze():
cfg.module_enable_status.nlm = True
cfg.nlm.h = 15
cfg.ceh.clip_limit = 0.04
cfg.eeh.flat_threshold = 1
cfg.eeh.flat_threshold = 2
cfg.eeh.edge_gain = 1280
cfg.hsc.saturation_gain = 320
cfg.bcc.contrast_gain = 280
pipeline = Pipeline(cfg)
pgm_path = 'raw/color_checker.pgm'
bayer = skimage.io.imread(pgm_path).astype(np.uint16)
_, intermediates = pipeline.execute(bayer, save_intermediates=True)
for module_name, result in intermediates.items():
output = pipeline.get_output(result)
output_path = op.join(OUTPUT_DIR, '{}.jpg'.format(module_name))
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
cv2.imwrite(output_path, output)
if __name__ == '__main__':
print('Processing test raw...')
demo_test_raw()
print('Processing Nikon D3x raw...')
demo_nikon_d3x()
print('Processing Nikon D3x raw with intermediate results...')
demo_intermediate_results()