forked from Ryanshuai/BM3D_py
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathbm3d_1st_step.py
111 lines (89 loc) · 4.25 KB
/
bm3d_1st_step.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
108
109
110
111
import numpy as np
import cv2
from utils import ind_initialize, get_kaiserWindow, sd_weighting
from precompute_BM import precompute_BM
from bior_2d import bior_2d_forward, bior_2d_reverse
from dct_2d import dct_2d_forward, dct_2d_reverse
from image_to_patches import image2patches
from build_3D_group import build_3D_group
from ht_filtering_hadamard import ht_filtering_hadamard
def bm3d_1st_step(sigma, img_noisy, nHard, kHard, NHard, pHard, lambdaHard3D, tauMatch, useSD, tau_2D):
height, width = img_noisy.shape[0], img_noisy.shape[1]
row_ind = ind_initialize(height - kHard + 1, nHard, pHard)
column_ind = ind_initialize(width - kHard + 1, nHard, pHard)
kaiserWindow = get_kaiserWindow(kHard)
ri_rj_N__ni_nj, threshold_count = precompute_BM(img_noisy, kHW=kHard, NHW=NHard, nHW=nHard, tauMatch=tauMatch)
group_len = int(np.sum(threshold_count))
group_3D_table = np.zeros((group_len, kHard, kHard))
weight_table = np.zeros((height, width))
all_patches = image2patches(img_noisy, kHard, kHard) # i_j_ipatch_jpatch__v
if tau_2D == 'DCT':
fre_all_patches = dct_2d_forward(all_patches)
else: # 'BIOR'
fre_all_patches = bior_2d_forward(all_patches)
acc_pointer = 0
for i_r in row_ind:
for j_r in column_ind:
nSx_r = threshold_count[i_r, j_r]
group_3D = build_3D_group(fre_all_patches, ri_rj_N__ni_nj[i_r, j_r], nSx_r)
group_3D, weight = ht_filtering_hadamard(group_3D, sigma, lambdaHard3D, not useSD)
group_3D = group_3D.transpose((2, 0, 1))
group_3D_table[acc_pointer:acc_pointer + nSx_r] = group_3D
acc_pointer += nSx_r
if useSD:
weight = sd_weighting(group_3D)
weight_table[i_r, j_r] = weight
if tau_2D == 'DCT':
group_3D_table = dct_2d_reverse(group_3D_table)
else: # 'BIOR'
group_3D_table = bior_2d_reverse(group_3D_table)
# group_3D_table = np.maximum(group_3D_table, 0)
# for i in range(1000):
# patch = group_3D_table[i]
# print(i, '----------------------------')
# print(patch)
# print(np.min(patch))
# print(np.max(patch))
# print(np.sum(patch))
# cv2.imshow('', patch.astype(np.uint8))
# cv2.waitKey()
# aggregation part
numerator = np.zeros_like(img_noisy, dtype=np.float64)
denominator = np.zeros((img_noisy.shape[0] - 2 * nHard, img_noisy.shape[1] - 2 * nHard), dtype=np.float64)
denominator = np.pad(denominator, nHard, 'constant', constant_values=1.)
acc_pointer = 0
for i_r in row_ind:
for j_r in column_ind:
nSx_r = threshold_count[i_r, j_r]
N_ni_nj = ri_rj_N__ni_nj[i_r, j_r]
group_3D = group_3D_table[acc_pointer:acc_pointer + nSx_r]
acc_pointer += nSx_r
weight = weight_table[i_r, j_r]
for n in range(nSx_r):
ni, nj = N_ni_nj[n]
patch = group_3D[n]
numerator[ni:ni + kHard, nj:nj + kHard] += patch * kaiserWindow * weight
denominator[ni:ni + kHard, nj:nj + kHard] += kaiserWindow * weight
img_basic = numerator / denominator
return img_basic
if __name__ == '__main__':
from utils import add_gaussian_noise, symetrize
# <hyper parameter> -------------------------------------------------------------------------------
sigma = 20
nHard = 16
kHard = 8
NHard = 16
pHard = 3
lambdaHard3D = 2.7 # ! Threshold for Hard Thresholding
tauMatchHard = 2500 if sigma < 35 else 5000 # ! threshold determinates similarity between patches
useSD_h = False
tau_2D_hard = 'BIOR'
# <\ hyper parameter> -----------------------------------------------------------------------------
img = cv2.imread('test_data/image/Cameraman.png', cv2.IMREAD_GRAYSCALE)
img_noisy = add_gaussian_noise(img, sigma)
# img_noisy = cv2.imread('matlab_officialfg_compare/noisy_image.png', cv2.IMREAD_GRAYSCALE)
img_noisy_p = symetrize(img_noisy, nHard)
img_basic = bm3d_1st_step(sigma, img_noisy_p, nHard, kHard, NHard, pHard, lambdaHard3D, tauMatchHard, useSD_h,
tau_2D_hard)
img_basic = img_basic[nHard: -nHard, nHard: -nHard]
# cv2.imwrite('y_basic.png', img_basic.astype(np.uint8))