diff --git a/docs/calib.md b/docs/calib.md new file mode 100644 index 0000000..fcdc972 --- /dev/null +++ b/docs/calib.md @@ -0,0 +1,17 @@ +# Noise calibration + +## Prerequisites + +### 1. Rawpy + +The calibration code the raspy package, which can be installed based on this instructions(https://pypi.org/project/rawpy/) + +Note: If rawpy cannot be installed on Mac(m1/m2), you can refer to this issue(https://github.com/letmaik/rawpy/issues/171) + +### 2. data + +It is recommended that you organize the data to be calibrated into separate folders for each camera. Within each camera folder, you should further divide the data into several subfolders based on different ISO values. + +## calibration process + +For the calibration process, you can execute all steps at once by directly following the code given in the main function, or you can perform each step separately according to your needs. These steps include selecting the positions of color blocks, calibrating the color blocks to obtain K, calibrating dark images to obtain other parameters, and fitting log(K) and log(variance). \ No newline at end of file diff --git a/tools/calib/calib_color.py b/tools/calib/calib_color.py new file mode 100644 index 0000000..3aaf4c3 --- /dev/null +++ b/tools/calib/calib_color.py @@ -0,0 +1,135 @@ +import numpy as np +import rawpy +from utils import * +import matplotlib.pyplot as plt +import json + +# Fit K and Var(No) using the formula Var(D) = K(KI) + Var(No) +def calib_color_per_iso(cam_dir, iso): + cur_iso_path = os.path.join(os.path.join(cam_dir,'color'), str(iso)) + raw_imgs = os.listdir(cur_iso_path) + if '.DS_Store' in raw_imgs: + raw_imgs.remove('.DS_Store') + + # Retrieve previously stored position information + camera = cam_dir.split('/')[-1] + block_position_path = f'./pos/{camera}_block_pos.npy' + block_positoins = np.load(block_position_path) + KI = np.zeros((4,len(block_positoins))) + var_D = np.zeros((4,len(block_positoins))) + for raw_img in raw_imgs: + cur_raw_path = os.path.join(cur_iso_path, raw_img) + raw = rawpy.imread(cur_raw_path) + black_level = raw.black_level_per_channel + white_point = raw.camera_white_level_per_channel + + raw_vis = raw.raw_image_visible.copy() + raw_pattern = raw.raw_pattern + raw = raw.raw_image_visible.astype(np.float32) + # rggb_img = bayer2_Rggb(raw) + rggb_img = pack_raw_bayer(raw_vis, raw_pattern) + rggb_img = rggb_img.transpose(1,2,0).astype(np.int64) + + rggb_img -= black_level + # rggb_img = rggb_img.astype(np.float32) + # rggb_img /= (np.array(white_point) - np.array(black_level)) + + for i, pos in enumerate(block_positoins): + minx,miny,w,h = pos + maxx,maxy = minx+w, miny+h + minx //= 2 + miny //= 2 + maxx //= 2 + maxy //= 2 + KI[:,i] += np.mean(rggb_img[miny:maxy,minx:maxx,:],axis=(0,1)) + # plt.imshow(rggb_img[miny:maxy,minx:maxx,:].mean(-1)) + # print(rggb_img[miny:maxy,minx:maxx,:].max()) + # print(rggb_img[miny:maxy,minx:maxx,:].min()) + # plt.show() + var_D[:,i] += np.var(rggb_img[miny:maxy,minx:maxx,:],axis=(0,1)) + + KI /= len(raw_imgs) + + var_D /= len(raw_imgs) + K, var_No = np.zeros((4)), np.zeros((4)) + for i in range(4): + K[i], var_No[i] = linear_regression(KI[i], var_D[i]) + print(iso) + # plt.scatter(KI[1],var_D[1]) + # plt.show() + for i in range(4): + plt.scatter(KI[i],var_D[i]) + fig_dir = './figs/' + if not os.path.exists(fig_dir): + os.makedirs(fig_dir) + fig_name = fig_dir + camera + '/' + f'{iso}_K{i}.png' + plt.savefig(fig_name) + plt.close() + return K,var_No + + +def calib_color_per_iso_whole(cam_dir, iso): + cur_iso_path = os.path.join(os.path.join(cam_dir,'color'), str(iso)) + raw_imgs = os.listdir(cur_iso_path) + if '.DS_Store' in raw_imgs: + raw_imgs.remove('.DS_Store') + + example_img_path = os.path.join(cur_iso_path,raw_imgs[0]) + rgb_img_example = raw2rgb(example_img_path) + rect = select_whole_positions(rgb_img_example) + # rect = (1514, 1018, 2543, 1897) + minx,miny,w,h = rect + maxx,maxy = minx+w, miny+h + minx //= 2 + miny //= 2 + maxx //= 2 + maxy //= 2 + print(rect) + total_imgs = np.zeros((len(raw_imgs),maxy-miny,maxx-minx,4)) + for i,raw_img in enumerate(raw_imgs): + cur_raw_path = os.path.join(cur_iso_path, raw_img) + raw = rawpy.imread(cur_raw_path) + black_level = raw.black_level_per_channel + white_point = raw.camera_white_level_per_channel + + raw_vis = raw.raw_image_visible.copy() + raw_pattern = raw.raw_pattern + raw = raw.raw_image_visible.astype(np.float32) + # rggb_img = bayer2_Rggb(raw) + rggb_img = pack_raw_bayer(raw_vis, raw_pattern) + rggb_img = rggb_img.transpose(1,2,0).astype(np.int64) + + rggb_img -= black_level + rggb_img = rggb_img.astype(np.float32) + rggb_img /= (np.array(white_point) - np.array(black_level)) + + total_imgs[i] = rggb_img[miny:maxy,minx:maxx,:] + + mean_img = np.mean(total_imgs,axis=0) + var_img = np.var(total_imgs,axis=0) + + # K, var_No = np.zeros((4)), np.zeros((4)) + # for i in range(4): + # K[i], var_No[i] = linear_regression(mean_img, var_D[i]) + plt.scatter(mean_img.reshape(-1,4)[:,1],var_img.reshape(-1,4)[:,1]) + plt.show() + # return K,var_No + +def calib_color_per_camera(cam_dir): + color_dir = os.path.join(cam_dir, 'color') + iso_list = sorted(os.listdir(color_dir)) + if '.DS_Store' in iso_list: + iso_list.remove('.DS_Store') + color_calib_params = dict() + color_calib_params['K_list'], color_calib_params['var_No_list'] = dict(), dict() + for iso in iso_list: + K, var_No = calib_color_per_iso(cam_dir,iso=int(iso)) + color_calib_params['K_list'][iso] = K.tolist() + color_calib_params['var_No_list'][iso] = var_No.tolist() + + cam_name = cam_dir.split('/')[-1] + color_calib_params_dir = os.path.join('./result/calib/color_calib_params',cam_name) + if not os.path.exists(color_calib_params_dir): + os.makedirs(color_calib_params_dir) + with open(os.path.join(color_calib_params_dir,'color_calib_params.json'),'w') as json_file: + json.dump(color_calib_params,json_file) \ No newline at end of file diff --git a/tools/calib/calib_dark.py b/tools/calib/calib_dark.py new file mode 100644 index 0000000..9d4bbc5 --- /dev/null +++ b/tools/calib/calib_dark.py @@ -0,0 +1,140 @@ +import numpy as np +import rawpy +from utils import * +import matplotlib.pyplot as plt +import json +import scipy + +def calib_dark_per_iso(cam_dir,iso): + cur_iso_path = os.path.join(os.path.join(cam_dir,'dark'), str(iso)) + raw_imgs = os.listdir(cur_iso_path) + if '.DS_Store' in raw_imgs: + raw_imgs.remove('.DS_Store') + + r = 400 + # Read + sigma_read = np.zeros((4, len(raw_imgs)), dtype=np.float32) + mean_read = np.zeros((4, len(raw_imgs)), dtype=np.float32) + r2_read = np.zeros((4, len(raw_imgs)), dtype=np.float32) + # row + sigma_row = np.zeros(len(raw_imgs), dtype=np.float32) + mean_row = np.zeros(len(raw_imgs), dtype=np.float32) + r2_row = np.zeros(len(raw_imgs), dtype=np.float32) + # TL + sigma_TL = np.zeros((4, len(raw_imgs)), dtype=np.float32) + mean_TL = np.zeros((4, len(raw_imgs)), dtype=np.float32) + r2_TL = np.zeros((4, len(raw_imgs)), dtype=np.float32) + lamda = np.zeros((4, len(raw_imgs)), dtype=np.float32) + # Gauss + sigma_gauss = np.zeros((4, len(raw_imgs)), dtype=np.float32) + mean_gauss = np.zeros((4, len(raw_imgs)), dtype=np.float32) + r2_gauss = np.zeros((4, len(raw_imgs)), dtype=np.float32) + + for i, raw_img in enumerate(raw_imgs): + # print(i) + cur_raw_path = os.path.join(cur_iso_path, raw_img) + raw = rawpy.imread(cur_raw_path) + black_level = raw.black_level_per_channel + + # print(f'bl:{black_level}') + + raw_vis = raw.raw_image_visible.copy() + raw_pattern = raw.raw_pattern + raw = raw.raw_image_visible.astype(np.float32) + # rggb_img = bayer2_Rggb(raw) + + # Calculate R before converting to 4 channels + raw -= np.mean(black_level) + row_all = np.mean(raw[raw.shape[0]//2-r*2:raw.shape[0]//2+r*2,raw.shape[1]//2-r*2:raw.shape[1]//2+r*2], axis=1) + _, (sig_row,u_row,r_row) = scipy.stats.probplot(row_all,rvalue=True) + sigma_row[i] = sig_row + mean_row[i] = u_row + r2_row[i] = r_row**2 + + # Convert the image into RGGB four channels + rggb_img = pack_raw_bayer(raw_vis, raw_pattern) + rggb_img = rggb_img.transpose(1,2,0).astype(np.int64) + # Subtract the black level + rggb_img -= black_level + + # Crop out a square with a side length of 800 + H, W = rggb_img.shape[:2] + rggb_img = rggb_img[H//2-r:H//2+r, W//2-r:W//2+r, :] + + + # Iterate over the 4 channels + for c in range(4): + cur_channel_img = rggb_img[:, :, c] + + # Calculate the variance of TL (or the variance of Gaussian) + the variance of row, here recorded as the variance of read + _, (sig_read,u_read,r_read) = scipy.stats.probplot(cur_channel_img.reshape(-1),rvalue=True) + sigma_read[c][i] = sig_read + mean_read[c][i] = u_read + r2_read[c][i] = r_read**2 + + # Calculate TL + row_all_cur_channel = np.mean(cur_channel_img,axis=1) + cur_channel_img = cur_channel_img.astype(np.float32) + cur_channel_img -= row_all_cur_channel.reshape(-1,1) + X = cur_channel_img.reshape(-1) + lam = scipy.stats.ppcc_max(X) + lamda[c][i] = lam + _, (sig_TL,u_TL,r_TL) = scipy.stats.probplot(X,dist=scipy.stats.tukeylambda(lam), rvalue=True) + sigma_TL[c][i] = sig_TL + mean_TL[c][i] = u_TL + r2_TL[c][i] = r_TL**2 + + # Calculate gauss + _,(sig_gauss,u_gauss,r_gauss) = scipy.stats.probplot(X,rvalue=True) + sigma_gauss[c][i] = sig_gauss + mean_gauss[c][i] = u_gauss + r2_gauss[c][i] = r_gauss**2 + param = { + 'black_level':black_level, + 'lam':lamda.tolist(), + 'sigmaRead':sigma_read.tolist(), 'meanRead':mean_read.tolist(), 'r2Gs':r2_read.tolist(), + 'sigmaR':sigma_row.tolist(), 'meanR':mean_row.tolist(), 'r2R':r2_row.tolist(), + 'sigmaTL':sigma_TL.tolist(), 'meanTL':mean_TL.tolist(), 'r2TL':r2_TL.tolist(), + 'sigmaGs':sigma_gauss.tolist(), 'meanGs':mean_gauss.tolist(), 'r2Gs':r2_gauss.tolist(), + } + + param_channel_mean = { + 'black_level':black_level, + 'lam':np.mean(lamda,axis=0).tolist(), + 'sigmaRead':np.mean(sigma_read,axis=0).tolist(), 'meanRead':np.mean(mean_read,axis=0).tolist(), 'r2Gs':np.mean(r2_read,axis=0).tolist(), + 'sigmaR':sigma_row.tolist(), 'meanR':mean_row.tolist(), 'r2R':r2_row.tolist(), + 'sigmaTL':np.mean(sigma_TL,axis=0).tolist(), 'meanTL':np.mean(mean_TL,axis=0).tolist(), 'r2TL':np.mean(r2_TL,axis=0).tolist(), + 'sigmaGs':np.mean(sigma_gauss,axis=0).tolist(), 'meanGs':np.mean(mean_gauss,axis=0).tolist(), 'r2Gs':np.mean(r2_gauss,axis=0).tolist(), + } + + return param, param_channel_mean + + +# get noise params from dark imgs per camera +def calib_dark_per_camera(cam_dir): + dark_dir = os.path.join(cam_dir, 'dark') + iso_list = sorted(os.listdir(dark_dir)) + if '.DS_Store' in iso_list: + iso_list.remove('.DS_Store') + # if '400' in iso_list: + # iso_list.remove('400') + dark_calib_params = dict() + dark_calib_params_channel_mean = dict() + for iso in iso_list: + print(iso) + param, param_channel_mean = calib_dark_per_iso(cam_dir,iso=int(iso)) + dark_calib_params[iso] = param + dark_calib_params_channel_mean[iso] = param_channel_mean + + + cam_name = cam_dir.split('/')[-1] + if not os.path.exists('./result/calib/dark_calib_params'): + os.mkdir('./result/calib/dark_calib_params') + dark_calib_params_dir = os.path.join('./result/calib/dark_calib_params',cam_name) + if not os.path.exists(dark_calib_params_dir): + os.mkdir(dark_calib_params_dir) + + with open(os.path.join(dark_calib_params_dir,'dark_calib_params.json'),'w') as json_file: + json.dump(dark_calib_params,json_file) + with open(os.path.join(dark_calib_params_dir,'dark_calib_params_channel_mean.json'),'w') as json_file: + json.dump(dark_calib_params_channel_mean,json_file) \ No newline at end of file diff --git a/tools/calib/noise_calib_main.py b/tools/calib/noise_calib_main.py new file mode 100644 index 0000000..5863b75 --- /dev/null +++ b/tools/calib/noise_calib_main.py @@ -0,0 +1,106 @@ +from utils import * +import rawpy +import os +import numpy as np +import matplotlib.pyplot as plt +import json +import scipy +from calib_color import * +from calib_dark import * + +# fit (log) var_TL, var_gauss, var_row and K +def fit_log(cam_name): + sig_gauss = [] + sig_row = [] + sig_TL = [] + K = [] + K_points = [] + with open(f'./result/calib/color_calib_params/{cam_name}/color_calib_params.json','r') as f: + data = json.load(f) + K_list, var_No_list = data['K_list'], data['var_No_list'] + with open(f'./result/calib/dark_calib_params/{cam_name}/dark_calib_params_channel_mean.json','r') as f: + dark_calib_params_channel_mean = json.load(f) + + for iso, param in dark_calib_params_channel_mean.items(): + + # for i in range(4): + cur_k = np.mean(K_list[iso]) + K.append(cur_k) + sig_row.append(np.mean(param['sigmaR'])) + sig_TL.append(np.mean(param['sigmaTL'])) + sig_gauss.append(np.mean(param['sigmaGs'])) + # K_points.append(K) + + fig = plt.figure(figsize=(20,8)) + + axsig_row = fig.add_subplot(1,3,1) + axsig_TL = fig.add_subplot(1,3,2) + axsig_gauss = fig.add_subplot(1,3,3) + + axsig_row.set_title('log(sig_row) - log(K)') + axsig_TL.set_title('log(sig_TL) - log(K)') + axsig_gauss.set_title('log(sig_gauss) - log(K)') + + axsig_row, data_row = regr_plot(K, sig_row, ax=axsig_row, c1='red', c2='orange', log=True, label=True) + axsig_TL, data_TL = regr_plot(K, sig_TL, ax=axsig_TL, c1='red', c2='orange', log=True, label=True) + axsig_gauss, data_gauss = regr_plot(K, sig_gauss, ax=axsig_gauss, c1='red', c2='orange', log=True, label=True) + + # plt.show() + params = dict() + params['Kmin'] = min(min(values) for values in K_list.values()) + params['Kmax'] = max(max(values) for values in K_list.values()) + params['Row'] = dict() + params['TL'] = dict() + params['Gauss'] = dict() + + params['Row']['slope'] = data_row['k'] + params['Row']['bias'] = data_row['b'] + params['Row']['std'] = data_row['sig'] + + params['TL']['slope'] = data_TL['k'] + params['TL']['bias'] = data_TL['b'] + params['TL']['std'] = data_TL['sig'] + + params['Gauss']['slope'] = data_gauss['k'] + params['Gauss']['bias'] = data_gauss['b'] + params['Gauss']['std'] = data_gauss['sig'] + + cam_param_dir = '.result/cam_log_params/' + cam_name + '.json' + if not os.path.exists(cam_param_dir): + os.mkdir(cam_param_dir) + with open(cam_param_dir,'w') as json_file: + json.dump(params,json_file) + if not os.path.exists('.result/log_figs/'): + os.mkdir('./result/log_figs/') + plt.savefig('./result/log_figs/' + cam_name + '.png') + + +if __name__ == "__main__": + # fit_log('6d2') + # calib_color_per_camera(cam_dir='/Users/hyx/Code/CV/raw_denoising/calib/r10') + # calib_dark() + # get_block_positoins(cam_dir= '/Users/hyx/Code/CV/raw_denoising/calib/550d') + # calib_color_per_iso_whole(cam_dir='/Users/hyx/Code/CV/raw_denoising/calib/6d2',iso=1600) + # calib_dark_per_iso('/Users/hyx/Code/CV/raw_denoising/calib/d5200',1600) + + # calib_dark_per_camera('/Users/hyx/Code/CV/raw_denoising/calib/d5200') + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--dir',type=str,default='./calib') + parser.add_argument('--mode',type=str) + directory = parser.parse_args().dir + cam_list = os.listdir(directory) + if '.DS_Store' in cam_list: + cam_list.remove('.DS_Store') + print(cam_list) + if parser.parse_args().mode == 'get_pos': + for cam in cam_list: + get_block_positoins(cam_dir=os.path.join(directory,cam)) + elif parser.parse_args().mode == 'calib': + for cam in cam_list: + cam_dir = os.path.join(directory,cam) + calib_dark_per_camera(cam_dir) + print('-------') + elif parser.parse_args().mode == 'fig_log': + for cam in cam_list: + fit_log(cam) \ No newline at end of file diff --git a/tools/calib/utils.py b/tools/calib/utils.py new file mode 100644 index 0000000..d694c0e --- /dev/null +++ b/tools/calib/utils.py @@ -0,0 +1,134 @@ +import numpy as np +import cv2 +import os +import rawpy +import imageio +from sklearn.linear_model import LinearRegression + +def pack_raw_bayer(raw: np.ndarray, raw_pattern: np.ndarray): + #pack Bayer image to 4 channels + R = np.where(raw_pattern==0) + G1 = np.where(raw_pattern==1) + B = np.where(raw_pattern==2) + G2 = np.where(raw_pattern==3) + + raw = raw.astype(np.uint16) + out = np.stack((raw[R[0][0]::2, R[1][0]::2], #RGBG + raw[G1[0][0]::2, G1[1][0]::2], + raw[B[0][0]::2, B[1][0]::2], + raw[G2[0][0]::2, G2[1][0]::2]), axis=0).astype(np.uint16) + + return out + + +# def bayer2rggb(bayer): +# H, W = bayer.shape +# return bayer.reshape(H//2, 2, W//2, 2).transpose(0, 2, 1, 3).reshape(H//2, W//2, 4) + +# read raw and return rgb +def raw2rgb(raw_path): + raw_image = rawpy.imread(raw_path) + rgb_image = raw_image.postprocess(use_camera_wb=True, half_size=False, no_auto_bright=True, output_bps=8) + return rgb_image + # imageio.imsave('output_rgb_image.jpg', rgb_image) + + + +# def load_rgb_image(rgb_root): +# rgb_path = os.path.join(rgb_root,"raw_visual") +# filenames = os.listdir(rgb_path) +# return cv2.imread(os.path.join(rgb_path,filenames[0])) + + +# Obtain and store the color block positions for a specific camera's color image +def get_block_positoins(cam_dir,iso=100): + raw_path = os.path.join(os.path.join(cam_dir,'color'),str(iso)) + raw_imgs = sorted(os.listdir(raw_path)) + if '.DS_Store' in raw_imgs: + raw_imgs.remove('.DS_Store') + example_raw_img = os.path.join(raw_path,raw_imgs[0]) + color_example_rgb = raw2rgb(example_raw_img) + block_positions = select_block_positions(color_example_rgb,calib_num=24) + cam_name = cam_dir.split('/')[-1] + save_file_name = f'{cam_name}_block_pos.npy' + if not os.path.exists('./pos'): + os.mkdir('./pos') + np.save(os.path.join('./pos',save_file_name),block_positions) + + +# select each block and ret the positions +def select_block_positions(rgb_image,calib_num=24): + positions = [] + for i in range(calib_num): + # rect = cv2.selectROI(rgb_image,False,False) + rect = cv2.selectROI(rgb_image,showCrosshair=True) + cv2.rectangle(rgb_image,rect,color=(23,128,62)) + print(i) + print(rect) + print('---------------') + positions.append(list(rect)) + + return np.array(positions) + +def select_whole_positions(rgb_image): + rect = cv2.selectROI(rgb_image,showCrosshair=True) + cv2.rectangle(rgb_image,rect,color=(23,128,62)) + + return rect + + +def linear_regression(x, y): + return np.polyfit(x, y, 1) + + +def regr_plot(x, y, log=True, ax=None, c1=None, c2=None, label=False): + x = np.array(x) + y = np.array(y) + if log: + x = np.log(x) + y = np.log(y) + ax.scatter(x, y) + + regr = LinearRegression() + regr.fit(x.reshape(-1,1), y) + a, b = float(regr.coef_), float(regr.intercept_) + # ax.set_title('log(sigR) | log(K)') + x_range = np.array([np.min(x), np.max(x)]) + std = np.mean((a*x+b-y)**2) ** 0.5 + + if c1 is not None: + if label: + label = f'k={a:.5f}, b={b:.5f}, std={std:.5f}' + ax.plot(x, regr.predict(x.reshape(-1,1)), linewidth = 2, c=c1, label=label) + else: + ax.plot(x, regr.predict(x.reshape(-1,1)), linewidth = 2, c=c1) + + if c2 is not None: + ax.plot(x_range, a*x_range+b+std, c=c2, linewidth = 1) + ax.plot(x_range, a*x_range+b-std, c=c2, linewidth = 1) + + data = {'k':a,'b':b,'sig':std} + + return ax, data + + +# raw2rgb('/Users/hyx/Code/CV/raw_denoising/calib/r10/color/100/IMG_0483.CR3') + +# data = np.load('/Users/hyx/Code/CV/raw_denoising/pos/r10_block_pos.npy') +# print(data) + +# block_positoins = np.load('/Users/hyx/Code/CV/raw_denoising/pos/r10_block_pos.npy') + +# raw_path = os.path.join(os.path.join('/Users/hyx/Code/CV/raw_denoising/calib/r10','color'),str(100)) +# raw_imgs = sorted(os.listdir(raw_path)) +# raw_imgs.remove('.DS_Store') +# example_raw_img = os.path.join(raw_path,raw_imgs[0]) +# color_example_rgb = raw2rgb(example_raw_img) + +# for i in range(24): +# # rect = cv2.selectROI(rgb_image,False,False) +# rect = block_positoins[i] +# cv2.rectangle(color_example_rgb,rect,color=(23,128,62)) +# cv2.imshow('example',color_example_rgb) +# cv2.waitKey(0) +