forked from alexandru-dinu/cae
-
Notifications
You must be signed in to change notification settings - Fork 0
/
smoothing.py
executable file
·60 lines (42 loc) · 1.62 KB
/
smoothing.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
import argparse
import numpy as np
from PIL import Image
from skimage.io import imsave
def lin_interp(n, p1, p2):
x = np.zeros((n, p1.shape[0], 128, 3))
for i in range(n):
a = (i + 1) / (n + 1)
x[i] = (1 - a) * p1 + a * p2
return x
def smooth(in_img, ws, out_img):
in_img = np.array(Image.open(in_img)) / 255.0
orig_img = in_img[24:-24, :1280, :]
in_img = in_img[:, 1280:, :]
# 6,10,128,128,3
patches = np.reshape(in_img, (6, 128, 10, 128, 3))
patches = np.transpose(patches, (0, 2, 1, 3, 4))
h = ws // 2
for i in range(5):
p1 = patches[i, :, 128 - h, :, :]
p2 = patches[i + 1, :, h, :, :]
x = lin_interp(ws, p1, p2)
patches[i, :, 128 - h:, :, :] = np.transpose(x[:h, :, :, :], (1, 0, 2, 3))
patches[i + 1, :, :h, :, :] = np.transpose(x[h:, :, :, :], (1, 0, 2, 3))
for j in range(9):
p3 = patches[:, j, :, 128 - h, :]
p4 = patches[:, j + 1, :, h, :]
x = lin_interp(ws, p3, p4)
patches[:, j, :, 128 - h:, :] = np.transpose(x[:h, :, :, :], (1, 2, 0, 3))
patches[:, j + 1, :, :h, :] = np.transpose(x[h:, :, :, :], (1, 2, 0, 3))
out = np.transpose(patches, (0, 2, 1, 3, 4))
out = np.reshape(out, (768, 1280, 3))
out = out[24:-24, :, :]
out = np.concatenate((orig_img, out), axis=1)
imsave(out_img, out)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--in_img', type=str, required=True)
parser.add_argument('--out_img', type=str, required=True)
parser.add_argument('--window_size', type=int, required=True)
args = parser.parse_args()
smooth(args.in_img, args.window_size, args.out_img)