-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
107 lines (84 loc) · 4.08 KB
/
main.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
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
import argparse
import os
import time
import numpy as np
from src.estimate import estimate_noise_curve
import src.utils as utils
parser = argparse.ArgumentParser(description='Video Signal-Dependent Noise Estimation via Inter-frame Prediction. '
'(c) 2022 Yanhao Li. Under license GNU AGPL.')
parser.add_argument('im_0', type=str,
help='First frame filename')
parser.add_argument('im_1', type=str,
help='Second frame filename')
parser.add_argument('-bins', type=int, default=16,
help='Number of bins (default: 16)')
parser.add_argument('-q', type=float, default=0.01,
help='Quantile of block pairs (default: 0.01)')
parser.add_argument('-s', type=int, default=5,
help='Search radius of patch matching (default: 5)')
parser.add_argument('-w', type=int, default=8,
help='Block size (default: 8)')
parser.add_argument('-T', type=int, default=-1,
help='Frequency separator (default: w+1)')
parser.add_argument('-th', type=int, default=3,
help='Thickness of ring for patch matching (default: 3)')
parser.add_argument('-g', default=False,
help='Whether the input images are in grayscale (default: False)', action='store_true')
parser.add_argument('-add_noise', default=False,
help='True for adding simulated noise, with noise model: variance = a + intensity * b. (default: False)',
action='store_true')
parser.add_argument('-noise_a', type=float, default=0.2,
help='Noise model parameter: a (default: 0.2)')
parser.add_argument('-noise_b', type=float, default=0.2,
help='Noise model parameter: b (default: 0.2)')
args = parser.parse_args()
def save_to_txt(intensities, variances, save_fname) -> None:
""" Save to a txt file of size (bins, channels * 2)
intensities: of size (channels, bins)
variances: of size (channels, bins)
save_fname: the txt filename
"""
assert intensities.shape == variances.shape
channels, bins = intensities.shape
out_data = np.zeros((bins, channels * 2))
for c in range(channels):
out_data[:, c*2] = intensities[c, :]
out_data[:, c*2+1] = variances[c, :]
len = int(np.log10(np.abs(out_data).max())) + 4
np.savetxt(save_fname, out_data, fmt=f'%{len}.3f')
def main():
print("Parameters:")
print(args)
print()
if args.T == -1:
args.T = args.w + 1
supported_ext = [".tif", ".tiff", ".dng"]
# verify the two images are in the same extension
_, extension_0 = os.path.splitext(args.im_0)
_, extension_1 = os.path.splitext(args.im_1)
assert extension_0 in supported_ext, \
f"Only `.tif`, `.tiff` and `.dng` formats are support, but `{extension_0}` was found."
assert extension_0 == extension_1, \
f"The two input images must be in the same format, but `{extension_0}` and `{extension_1}` were got."
img_0 = utils.read_img(args.im_0, grayscale=args.g)
img_1 = utils.read_img(args.im_1, grayscale=args.g)
if args.add_noise == True:
img_0 = utils.add_noise(img_0, args.noise_a, args.noise_b)
img_1 = utils.add_noise(img_1, args.noise_a, args.noise_b)
if img_0.shape != img_1.shape:
print("Error: The two input images should have the same size and the same channel")
quit()
if args.T > 2 * args.w - 3:
print("Error: Frequency separator T and block size w should satisfy T<=2*w-3")
quit()
start = time.time()
img_0 = img_0.astype(np.float32)
img_1 = img_1.astype(np.float32)
intensities, variances = estimate_noise_curve(img_0, img_1, w=args.w, T=args.T, th=args.th, q=args.q, bins=args.bins, s=args.s)
save_to_txt(intensities, variances, f"curve_s0.txt")
utils.plot_noise_curve(intensities, variances, fname=f"curve_s0.png")
print("The estimated noise curves are saved in `curve_s0.txt` and plotted in `curve_s0.png`")
print()
print(f"time spent: {time.time() - start} s")
if __name__ == "__main__":
main()