def bm3d_2nd_step(sigma, img_noisy, img_basic, nWien, kWien, NWien, pWien, tauMatch, useSD, tau_2D, lam): height, width = img_noisy.shape[0], img_noisy.shape[1] row_ind = ind_initialize(height - kWien + 1, nWien, pWien) column_ind = ind_initialize(width - kWien + 1, nWien, pWien) kaiserWindow = get_kaiserWindow(kWien) ri_rj_N__ni_nj, threshold_count = precompute_BM(img_basic, kHW=kWien, NHW=NWien, nHW=nWien, tauMatch=tauMatch) group_len = int(np.sum(threshold_count)) group_3D_table = np.zeros((group_len, kWien, kWien)) weight_table = np.ones((height, width)) noisy_patches = image2patches(img_noisy, k=kWien, p=pWien) # i_j_ipatch_jpatch__v basic_patches = image2patches(img_basic, k=kWien, p=pWien) # i_j_ipatch_jpatch__v if tau_2D == 'DCT': fre_noisy_patches = dct_2d_forward(noisy_patches) fre_basic_patches = dct_2d_forward(basic_patches) elif tau_2D == 'BIOR': # 'BIOR' fre_noisy_patches = bior_2d_forward(noisy_patches) fre_basic_patches = bior_2d_forward(basic_patches) else: fre_noisy_patches = np.zeros_like(noisy_patches) for i, patch in enumerate(noisy_patches): fre_noisy_patches[i] = ptv.tv1_2d(patch, lam) fre_basic_patches = np.zeros_like(basic_patches) for i, patch in enumerate(basic_patches): fre_basic_patches[i] = ptv.tv1_2d(patch, lam) fre_basic_patches = fre_basic_patches.reshape( (height - kWien + 1, height - kWien + 1, kWien, kWien)) 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_est = build_3D_group(fre_basic_patches, ri_rj_N__ni_nj[i_r, j_r], nSx_r) group_3D = group_3D_est 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) elif tau_2D == 'BIOR': # 'BIOR' group_3D_table = bior_2d_reverse(group_3D_table) else: pass # for i in range(1000): # patch = group_3D_table[i] # print(i, '----------------------------') # print(patch) # cv2.imshow('', patch.astype(np.uint8)) # cv2.waitKey() group_3D_table *= kaiserWindow numerator = np.zeros_like(img_noisy, dtype=np.float64) denominator = np.zeros_like(img_noisy, dtype=np.float64) acc_pointer = 0 for i_r in row_ind: for j_r in column_ind: if i_r == 264 and j_r == 264: print() 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 + kWien, nj:nj + kWien] += patch * weight denominator[ni:ni + kWien, nj:nj + kWien] += kaiserWindow * weight img_denoised = numerator / denominator return img_denoised
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() 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
def bm3d_2nd_step(sigma, img_noisy, img_basic, nWien, kWien, NWien, pWien, tauMatch, useSD, tau_2D): height, width = img_noisy.shape[0], img_noisy.shape[1] row_ind = ind_initialize(height - kWien + 1, nWien, pWien) column_ind = ind_initialize(width - kWien + 1, nWien, pWien) kaiserWindow = get_kaiserWindow(kWien) ri_rj_N__ni_nj, threshold_count = precompute_BM(img_basic, kHW=kWien, NHW=NWien, nHW=nWien, tauMatch=tauMatch) group_len = int(np.sum(threshold_count)) group_3D_table = np.zeros((group_len, kWien, kWien)) weight_table = np.zeros((height, width)) noisy_patches = image2patches(img_noisy, kWien, kWien) # i_j_ipatch_jpatch__v basic_patches = image2patches(img_basic, kWien, kWien) # i_j_ipatch_jpatch__v if tau_2D == 'DCT': fre_noisy_patches = dct_2d_forward(noisy_patches) fre_basic_patches = dct_2d_forward(basic_patches) else: # 'BIOR' fre_noisy_patches = bior_2d_forward(noisy_patches) fre_basic_patches = bior_2d_forward(basic_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_img = build_3D_group(fre_noisy_patches, ri_rj_N__ni_nj[i_r, j_r], nSx_r) group_3D_est = build_3D_group(fre_basic_patches, ri_rj_N__ni_nj[i_r, j_r], nSx_r) group_3D, weight = wiener_filtering_hadamard( group_3D_img, group_3D_est, sigma, 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) # for i in range(1000): # patch = group_3D_table[i] # print(i, '----------------------------') # print(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 * nWien, img_noisy.shape[1] - 2 * nWien), dtype=np.float64) denominator = np.pad(denominator, nWien, '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 + kWien, nj:nj + kWien] += patch * kaiserWindow * weight denominator[ni:ni + kWien, nj:nj + kWien] += kaiserWindow * weight img_denoised = numerator / denominator return img_denoised
def bm3d_1st_step_with_TV(sigma, img_noisy, nHard, kHard, NHard, pHard, lambdaHard3D, tauMatch, useSD, tau_2D, lamb): 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) for i in range(height): for j in range(width): assert ri_rj_N__ni_nj[i][j][0][0] == i assert ri_rj_N__ni_nj[i][j][0][1] == j 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, k=kHard, p=pHard) # i_j_ipatch_jpatch__v if tau_2D == 'DCT': fre_all_patches = dct_2d_forward(all_patches) elif tau_2D == 'BIOR': # 'BIOR' fre_all_patches = bior_2d_forward(all_patches) else: fre_all_patches = all_patches fre_all_patches = fre_all_patches.reshape((height - kHard + 1, height - kHard + 1, kHard, kHard)) 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_temp, weight = ht_filtering_hadamard(group_3D, sigma, lambdaHard3D, not useSD) group_3D = trendFilter3D(group_3D, lamb) 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) elif tau_2D == 'BIOR': # 'BIOR' group_3D_table = bior_2d_reverse(group_3D_table) else: group_3D_table *= kaiserWindow 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] if tau_2D == 'DCT' or tau_2D == 'BIOR': numerator[ni:ni + kHard, nj:nj + kHard] += patch * kaiserWindow * weight denominator[ni:ni + kHard, nj:nj + kHard] += kaiserWindow * weight elif tau_2D == 'TV': numerator[ni:ni + kHard, nj:nj + kHard] += patch * weight denominator[ni:ni + kHard, nj:nj + kHard] += kaiserWindow * weight img_basic = numerator / denominator return img_basic
def bm3d_1st_step(sigma, img_noisy, nHard, kHard, NHard, pHard, useSD, tau_2D): height, width = img_noisy.shape[0], img_noisy.shape[1] lambdaHard3D = 2.7 # ! Threshold for Hard Thresholding tauMatch = 3 * (2500 if sigma < 35 else 5000 ) # ! threshold determinates similarity between patches 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, k=kHard, p=pHard) # 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) fre_all_patches = fre_all_patches.reshape( (height - kHard + 1, height - kHard + 1, kHard, kHard)) 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 = group_3D.reshape(kHard * kHard, nSx_r) # group_3D, weight = ht_filtering_hadamard(group_3D, sigma, lambdaHard3D, not useSD) group_3D = group_3D.reshape(kHard, kHard, nSx_r) 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 weight_table[i_r, j_r] = 1 if tau_2D == 'DCT': group_3D_table = dct_2d_reverse(group_3D_table) else: # 'BIOR' group_3D_table = bior_2d_reverse(group_3D_table) # for i in range(1000): # patch = group_3D_table[i] # print(i, '----------------------------') # print(patch) # cv2.imshow('', patch.astype(np.uint8)) # cv2.waitKey() group_3D_table *= kaiserWindow numerator = np.zeros_like(img_noisy, dtype=np.float) denominator = np.zeros_like(img_noisy, dtype=np.float) 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 * weight denominator[ni:ni + kHard, nj:nj + kHard] += kaiserWindow * weight img_basic = numerator / denominator img_basic = img_basic.astype(np.uint8) return img_basic
def bm3d_1st_step_block_mean(sigma, img_noisy, nHard, kHard, NHard, pHard, lambdaHard3D, tauMatch, useSD, tau_2D, block_mean=False): 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, k=kHard, p=pHard) # 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) fre_all_patches = fre_all_patches.reshape((height - kHard + 1, height - kHard + 1, kHard, kHard)) 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) if block_mean: group_3D, weight = get_mean(group_3D) else: 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) numerator = np.zeros_like(img_noisy, dtype=np.float64) denominator = np.zeros_like(img_noisy, dtype=np.float64) 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