-
Notifications
You must be signed in to change notification settings - Fork 13
/
demo.py
74 lines (65 loc) · 2.26 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
# -*- coding: utf-8 -*-
import os
import time
import cv2
import numpy as np
from python_color_transfer.color_transfer import ColorTransfer
def demo():
cur_dir = os.path.abspath(os.path.dirname(__file__))
img_folder = os.path.join(cur_dir, "imgs")
img_names = [
"scotland_house.png",
"house.jpeg",
"fallingwater.png",
"tower.jpeg",
]
ref_names = [
"scotland_plain.png",
"hats.png",
"autumn.jpg",
"sunset.jpg",
]
out_names = [
"scotland_display.png",
"house_display.png",
"fallingwater_display.png",
"tower_display.png",
]
img_paths = [os.path.join(img_folder, x) for x in img_names]
ref_paths = [os.path.join(img_folder, x) for x in ref_names]
out_paths = [os.path.join(img_folder, x) for x in out_names]
# cls init
PT = ColorTransfer()
for img_path, ref_path, out_path in zip(img_paths, ref_paths, out_paths):
# read input img
img_arr_in = cv2.imread(img_path)
[h, w, c] = img_arr_in.shape
print(f"{img_path}: {h}x{w}x{c}")
# read reference img
img_arr_ref = cv2.imread(ref_path)
[h, w, c] = img_arr_ref.shape
print(f"{ref_path}: {h}x{w}x{c}")
# pdf transfer
t0 = time.time()
img_arr_reg = PT.pdf_transfer(img_arr_in=img_arr_in,
img_arr_ref=img_arr_ref,
regrain=True)
print(f"Pdf transfer time: {time.time() - t0:.2f}s")
# mean transfer
t0 = time.time()
img_arr_mt = PT.mean_std_transfer(img_arr_in=img_arr_in,
img_arr_ref=img_arr_ref)
print(f"Mean std transfer time: {time.time() - t0:.2f}s")
# lab transfer
t0 = time.time()
img_arr_lt = PT.lab_transfer(img_arr_in=img_arr_in,
img_arr_ref=img_arr_ref)
print(f"Lab mean std transfer time: {time.time() - t0:.2f}s")
# display
img_arr_out = np.concatenate(
(img_arr_in, img_arr_ref, img_arr_mt, img_arr_lt, img_arr_reg),
axis=1)
cv2.imwrite(out_path, img_arr_out)
print(f"Saved to {out_path}\n")
if __name__ == "__main__":
demo()