"""Forward processing of raw data to sRGB images. Unprocessing Images for Learned Raw Denoising http://timothybrooks.com/tech/unprocessing """ import rawpy import rawpy.enhance import exifread import numpy as np import torch import torch.nn as nn import torch.distributions as tdist from scipy import stats from utils import * import random from .unprocess import random_gains Dual_ISO_Cameras = ['SonyA7S2'] HALF_CLIP = 2 def bpc_aug(data, ratio=1e-6, wp=1): B, C, H, W = data.shape npoints = int(2 * np.random.rand() * ratio * data.numel()) for i in range(npoints): x = np.random.randint(H) y = np.random.randint(W) c = np.random.randint(C) b = np.random.randint(B) data[b, c, x, y] = wp return data def data_aug(data, choice, bias=0, rot=False): if choice[0] == 1: data = np.flip(data, axis=2+bias) if choice[1] == 1: data = np.flip(data, axis=3+bias) return data def inverse_VST_torch(x, noiseparam, iso_list, wp=1): x = x * wp b = len(iso_list) for i in range(b): iso = iso_list[i].item() gain = noiseparam[iso]['Kmax'] sigma = noiseparam[iso]['sigGs'] x[i] = (x[i] / 2.0)**2 - 3.0/8.0 - sigma**2 / gain**2 x[i] = x[i] * gain x = x / wp return x def pack_raw_bayer(raw, wp=1023, clip=True): #pack Bayer image to 4 channels im = raw.raw_image_visible.astype(np.float32) raw_pattern = raw.raw_pattern R = np.where(raw_pattern==0) G1 = np.where(raw_pattern==1) B = np.where(raw_pattern==2) G2 = np.where(raw_pattern==3) white_point = wp img_shape = im.shape H = img_shape[0] W = img_shape[1] out = np.stack((im[R[0][0]:H:2, R[1][0]:W:2], #RGBG im[G1[0][0]:H:2, G1[1][0]:W:2], im[B[0][0]:H:2, B[1][0]:W:2], im[G2[0][0]:H:2, G2[1][0]:W:2]), axis=0).astype(np.float32) black_level = np.array(raw.black_level_per_channel)[:,None,None].astype(np.float32) out = (out - black_level) / (white_point - black_level) out = np.clip(out, 0.0, 1.0) if clip else out return out def postprocess_bayer(rawpath, img4c, white_point=1023): if torch.is_tensor(img4c): img4c = img4c.detach() img4c = img4c[0].cpu().float().numpy() img4c = np.clip(img4c, 0, 1) #unpack 4 channels to Bayer image raw = rawpy.imread(rawpath) raw_pattern = raw.raw_pattern R = np.where(raw_pattern==0) G1 = np.where(raw_pattern==1) G2 = np.where(raw_pattern==3) B = np.where(raw_pattern==2) black_level = np.array(raw.black_level_per_channel)[:,None,None] img4c = img4c * (white_point - black_level) + black_level img_shape = raw.raw_image_visible.shape H = img_shape[0] W = img_shape[1] raw.raw_image_visible[R[0][0]:H:2, R[1][0]:W:2] = img4c[0, :,:] raw.raw_image_visible[G1[0][0]:H:2,G1[1][0]:W:2] = img4c[1, :,:] raw.raw_image_visible[B[0][0]:H:2,B[1][0]:W:2] = img4c[2, :,:] raw.raw_image_visible[G2[0][0]:H:2,G2[1][0]:W:2] = img4c[3, :,:] # out = raw.postprocess(use_camera_wb=False, user_wb=[1,1,1,1], half_size=True, no_auto_bright=True, output_bps=8, bright=1, user_black=None, user_sat=None) # out = raw.postprocess(use_camera_wb=False, user_wb=[1.96875, 1, 1.444, 1], half_size=True, no_auto_bright=True, output_bps=8, bright=1, user_black=None, user_sat=None) out = raw.postprocess(use_camera_wb=True, half_size=False, no_auto_bright=True, output_bps=8, bright=1, user_black=None, user_sat=None) return out def postprocess_bayer_v2(rawpath, img4c): with rawpy.imread(rawpath) as raw: out_srgb = raw2rgb_postprocess(img4c.detach(), raw) return out_srgb def apply_gains(bayer_images, wbs): """Applies white balance to a batch of Bayer images.""" N, C, _, _ = bayer_images.shape wbs = wbs.repeat((N,1)).view(N, C, 1, 1) outs = bayer_images * wbs return outs def apply_ccms(images, ccms): """Applies color correction matrices.""" images = images.permute( 0, 2, 3, 1) # Permute the image tensor to BxHxWxC format from BxCxHxW format images = images[:, :, :, None, :] ccms = ccms[:, None, None, :, :] outs = torch.sum(images * ccms, dim=-1) # Re-Permute the tensor back to BxCxHxW format outs = outs.permute(0, 3, 1, 2) return outs def gamma_compression(images, gamma=2.2): """Converts from linear to gamma space.""" outs = torch.clamp(images, min=1e-8) ** (1 / gamma) # outs = (1 + gamma[0]) * np.power(images, 1.0/gamma[1]) - gamma[0] + gamma[2]*images outs = torch.clamp((outs*255).int(), min=0, max=255).float() / 255 return outs def raw2LRGB(bayer_images): """RGBG -> linear RGB""" lin_rgb = torch.stack([ bayer_images[:,0,...], torch.mean(bayer_images[:, [1,3], ...], dim=1), bayer_images[:,2,...]], dim=1) return lin_rgb def process(bayer_images, wbs, cam2rgbs, gamma=2.2): """Processes a batch of Bayer RGBG images into sRGB images.""" # White balance. bayer_images = apply_gains(bayer_images, wbs) # Binning bayer_images = torch.clamp(bayer_images, min=0.0, max=1.0) images = raw2LRGB(bayer_images) # Color correction. images = apply_ccms(images, cam2rgbs) # Gamma compression. images = torch.clamp(images, min=0.0, max=1.0) images = gamma_compression(images, gamma) return images def raw2rgb(packed_raw, raw): """Raw2RGB pipeline (preprocess version)""" wb = np.array(raw.camera_whitebalance) wb /= wb[1] cam2rgb = raw.color_matrix[:3, :3] if cam2rgb[0,0] == 0: cam2rgb = np.eye(3, dtype=np.float32) if isinstance(packed_raw, np.ndarray): packed_raw = torch.from_numpy(packed_raw).float() wb = torch.from_numpy(wb).float().to(packed_raw.device) cam2rgb = torch.from_numpy(cam2rgb).float().to(packed_raw.device) out = process(packed_raw[None], wbs=wb[None], cam2rgbs=cam2rgb[None], gamma=2.2)[0, ...].numpy() return out def raw2rgb_v2(packed_raw, wb, ccm): if torch.is_tensor(packed_raw): packed_raw = packed_raw.detach().cpu().float() else: packed_raw = torch.from_numpy(packed_raw).float() wb = torch.from_numpy(wb).float() cam2rgb = torch.from_numpy(ccm).float() out = process(packed_raw[None], wbs=wb[None], cam2rgbs=cam2rgb[None], gamma=2.2)[0, ...].numpy() return out.transpose(1,2,0) def raw2rgb_postprocess(packed_raw, raw): """Raw2RGB pipeline (postprocess version)""" assert packed_raw.ndimension() == 4 wb = np.array(raw.camera_whitebalance) wb /= wb[1] cam2rgb = raw.color_matrix[:3, :3] if cam2rgb[0,0] == 0: cam2rgb = np.eye(3, dtype=np.float32) wb = torch.from_numpy(wb[None]).float().to(packed_raw.device) cam2rgb = torch.from_numpy(cam2rgb[None]).float().to(packed_raw.device) out = process(packed_raw, wbs=wb, cam2rgbs=cam2rgb, gamma=2.2) # out = raw.postprocess(use_camera_wb=True, half_size=False, no_auto_bright=True, output_bps=16) return out # def get_specific_noise_params(camera_type=None, iso='100'): # cam_noisy_params = {} # cam_noisy_params['IMX686'] = { # '100':{'K':0.1366021, 'sigGs':0.6926457, 'sigGssig':0.002096}, # '6400':{'K':8.7425333, 'sigGs':14.303619546153575, 'sigGssig':0.0696716845864088}, # } # if camera_type in cam_noisy_params: # return cam_noisy_params[camera_type][iso] # else: # log(f'''Warning: we have not test the noisy parameters of camera "{camera_type}". Now we use NikonD850's parameters to test.''') # return cam_noisy_params['IMX686'] def get_camera_noisy_params(camera_type=None): cam_noisy_params = {} cam_noisy_params['OS04J10'] = {#K_ISO=(0.01134,0.56113) 'K_ISO': (0.01143, 0.4127), 'ISOmin':100, 'ISOmax':1593, 'tmin':30, 'tmax':30, # exposure time(ms) 'Kmin':0.44193, 'Kmax':2.92427, 'lam':-0.094, 'q':2.442e-04, 'wp':4095, 'bl':256, 'sigRk':0.81696, 'sigRb':-2.69453, 'sigRsig':0.04015, 'sigTLk':0.82984, 'sigTLb':-0.76241, 'sigTLsig':0.06671, 'sigGsk':0.89423, 'sigGsb':-0.08060, 'sigGssig':0.04103, 'sigReadk':0.00772, 'sigReadb':0.65775, 'sigReadsig':0.34109, 'uReadk':0.00001, 'uReadb':0.00207, 'uReadsig':0.00059 } cam_noisy_params['SC850AI'] = { 'Kmin':1.42030, 'Kmax':1.70373, 'lam':-0.035, 'q':1/959, 'wp':1023, 'bl':64, 'sigRk':0.85792, 'sigRb':-2.82338, 'sigRsig':0.01944, 'sigTLk':-0.15504, 'sigTLb':0.71196, 'sigTLsig':0.04500, 'sigGsk':0.60800, 'sigGsb':0.17009, 'sigGssig':0.02261, 'sigReadk':0.60940, 'sigReadb':0.16932, 'sigReadsig':0.02257, 'uReadk':0.00000, 'uReadb':-0.00292, 'uReadsig':0.00400 } cam_noisy_params['SC450AI'] = { # ISO_3600_6052 (0.00065631, 0.02588613) 'Kmin':0.87071, 'Kmax':1.38576, 'lam':0.042, 'q':1/959, 'wp':1023, 'bl':64, 'sigRk':0.98972, 'sigRb':-2.69256, 'sigRsig':0.03678, 'sigTLk':0.75890, 'sigTLb':-0.51098, 'sigTLsig':0.04700, 'sigGsk':0.87135, 'sigGsb':-0.11635, 'sigGssig':0.04866, 'sigReadk':0.00044, 'sigReadb':0.30172, 'sigReadsig':0.12209, 'uReadk':-0.00005, 'uReadb':0.32066, 'uReadsig':0.01992 } # new calibration (smaller) cam_noisy_params['SC450'] = { # ISO_3600_6052 (0.00065631, 0.02588613) 'Kmin':0.87071, 'Kmax':1.38576, 'lam':0.042, 'q':1/959, 'wp':1023, 'bl':64, 'sigRk':0.97496, 'sigRb':-2.78752, 'sigRsig':0.02109, 'sigTLk':0.57693, 'sigTLb':-0.44032, 'sigTLsig':0.01515, 'sigGsk':0.88683, 'sigGsb':-0.26233, 'sigGssig':0.00695, 'sigReadk':0.00040, 'sigReadb':0.23393, 'sigReadsig':0.01396, 'uReadk':0.00000, 'uReadb':0.00076, 'uReadsig':0.00125 } cam_noisy_params['NikonD850'] = { 'Kmin':1.2, 'Kmax':2.4828, 'lam':-0.26, 'q':1/(2**14), 'wp':16383, 'bl':512, 'sigTLk':0.906, 'sigTLb':-0.6754, 'sigTLsig':0.035165, 'sigRk':0.8322, 'sigRb':-2.3326, 'sigRsig':0.301333, 'sigGsk':0.906, 'sigGsb':-0.1754, 'sigGssig':0.035165, } cam_noisy_params['IMX686'] = { # ISO-640~6400 'Kmin':-0.19118, 'Kmax':2.16820, 'lam':0.102, 'q':1/(2**10), 'wp':1023, 'bl':64, 'sigTLk':0.85187, 'sigTLb':0.07991, 'sigTLsig':0.02921, 'sigRk':0.87611, 'sigRb':-2.11455, 'sigRsig':0.03274, 'sigGsk':0.85187, 'sigGsb':0.67991, 'sigGssig':0.02921, } cam_noisy_params['SonyA7S2_lowISO'] = { 'Kmin':-1.67214, 'Kmax':0.42228, 'lam':-0.026, 'q':1/(2**14), 'wp':16383, 'bl':512, 'sigRk':0.78782, 'sigRb':-0.34227, 'sigRsig':0.02832, 'sigTLk':0.74043, 'sigTLb':0.86182, 'sigTLsig':0.00712, 'sigGsk':0.82966, 'sigGsb':1.49343, 'sigGssig':0.00359, 'sigReadk':0.82879, 'sigReadb':1.50601, 'sigReadsig':0.00362, 'uReadk':0.01472, 'uReadb':0.01129, 'uReadsig':0.00034, } cam_noisy_params['SonyA7S2_highISO'] = { 'Kmin':0.64567, 'Kmax':2.51606, 'lam':-0.025, 'q':1/(2**14), 'wp':16383, 'bl':512, 'sigRk':0.62945, 'sigRb':-1.51040, 'sigRsig':0.02609, 'sigTLk':0.74901, 'sigTLb':-0.12348, 'sigTLsig':0.00638, 'sigGsk':0.82878, 'sigGsb':0.44162, 'sigGssig':0.00153, 'sigReadk':0.82645, 'sigReadb':0.45061, 'sigReadsig':0.00156, 'uReadk':0.00385, 'uReadb':0.00674, 'uReadsig':0.00039, } cam_noisy_params['CRVD'] = { 'Kmin':1.31339, 'Kmax':3.95448, 'lam':0.015, 'q':1/(2**12), 'wp':4095, 'bl':240, 'sigRk':0.93368, 'sigRb':-2.19692, 'sigRsig':0.02473, 'sigGsk':0.95387, 'sigGsb':0.01552, 'sigGssig':0.00855, 'sigTLk':0.95495, 'sigTLb':0.01618, 'sigTLsig':0.00790 } if camera_type in cam_noisy_params: return cam_noisy_params[camera_type] else: log(f'''Warning: we have not test the noisy parameters of camera "{camera_type}". Now we use NikonD850's parameters to test.''') return cam_noisy_params['NikonD850'] def get_specific_noise_params(camera_type=None, iso='100'): iso = str(iso) cam_noisy_params = {} cam_noisy_params['SonyA7S2'] = { '50': {'Kmax': 0.047815, 'lam': 0.1474653, 'sigGs': 1.0164667, 'sigGssig': 0.005272454, 'sigTL': 0.70727646, 'sigTLsig': 0.004360543, 'sigR': 0.13997398, 'sigRsig': 0.0064381803, 'bias': 0, 'biassig': 0.010093017, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '64': {'Kmax': 0.0612032, 'lam': 0.13243394, 'sigGs': 1.0509665, 'sigGssig': 0.008081373, 'sigTL': 0.71535635, 'sigTLsig': 0.0056863446, 'sigR': 0.14346549, 'sigRsig': 0.006400559, 'bias': 0, 'biassig': 0.008690166, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '80': {'Kmax': 0.076504, 'lam': 0.1121489, 'sigGs': 1.180899, 'sigGssig': 0.011333668, 'sigTL': 0.7799473, 'sigTLsig': 0.009347968, 'sigR': 0.19540153, 'sigRsig': 0.008197397, 'bias': 0, 'biassig': 0.0107246125, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '100': {'Kmax': 0.09563, 'lam': 0.14875287, 'sigGs': 1.0067395, 'sigGssig': 0.0033682834, 'sigTL': 0.70181876, 'sigTLsig': 0.0037532174, 'sigR': 0.1391465, 'sigRsig': 0.006530218, 'bias': 0, 'biassig': 0.007235429, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '125': {'Kmax': 0.1195375, 'lam': 0.12904578, 'sigGs': 1.0279676, 'sigGssig': 0.007364685, 'sigTL': 0.6961967, 'sigTLsig': 0.0048687346, 'sigR': 0.14485553, 'sigRsig': 0.006731584, 'bias': 0, 'biassig': 0.008026363, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '160': {'Kmax': 0.153008, 'lam': 0.094135, 'sigGs': 1.1293099, 'sigGssig': 0.008340453, 'sigTL': 0.7258587, 'sigTLsig': 0.008032158, 'sigR': 0.19755602, 'sigRsig': 0.0082754735, 'bias': 0, 'biassig': 0.0101351, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '200': {'Kmax': 0.19126, 'lam': 0.07902429, 'sigGs': 1.2926387, 'sigGssig': 0.012171176, 'sigTL': 0.8117464, 'sigTLsig': 0.010250768, 'sigR': 0.22815849, 'sigRsig': 0.010726711, 'bias': 0, 'biassig': 0.011413908, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '250': {'Kmax': 0.239075, 'lam': 0.051688068, 'sigGs': 1.4345995, 'sigGssig': 0.01606571, 'sigTL': 0.8630922, 'sigTLsig': 0.013844714, 'sigR': 0.26271912, 'sigRsig': 0.0130637, 'bias': 0, 'biassig': 0.013569083, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '320': {'Kmax': 0.306016, 'lam': 0.040700804, 'sigGs': 1.7481371, 'sigGssig': 0.019626873, 'sigTL': 1.0334468, 'sigTLsig': 0.017629284, 'sigR': 0.3097104, 'sigRsig': 0.016202712, 'bias': 0, 'biassig': 0.017825918, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '400': {'Kmax': 0.38252, 'lam': 0.0222538, 'sigGs': 2.0595572, 'sigGssig': 0.024872316, 'sigTL': 1.1816813, 'sigTLsig': 0.02505812, 'sigR': 0.36209714, 'sigRsig': 0.01994737, 'bias': 0, 'biassig': 0.021005306, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '500': {'Kmax': 0.47815, 'lam': -0.0031342343, 'sigGs': 2.3956928, 'sigGssig': 0.030144656, 'sigTL': 1.31772, 'sigTLsig': 0.028629242, 'sigR': 0.42528257, 'sigRsig': 0.025104137, 'bias': 0, 'biassig': 0.02981831, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '640': {'Kmax': 0.612032, 'lam': 0.002566592, 'sigGs': 2.9662898, 'sigGssig': 0.045661453, 'sigTL': 1.6474211, 'sigTLsig': 0.04671843, 'sigR': 0.48839623, 'sigRsig': 0.031589635, 'bias': 0, 'biassig': 0.10000693, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '800': {'Kmax': 0.76504, 'lam': -0.008199721, 'sigGs': 3.5475867, 'sigGssig': 0.052318197, 'sigTL': 1.9346539, 'sigTLsig': 0.046128694, 'sigR': 0.5723769, 'sigRsig': 0.037824076, 'bias': 0, 'biassig': 0.025339302, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '1000': {'Kmax': 0.9563, 'lam': -0.021061005, 'sigGs': 4.2727833, 'sigGssig': 0.06972333, 'sigTL': 2.2795107, 'sigTLsig': 0.059203167, 'sigR': 0.6845563, 'sigRsig': 0.04879781, 'bias': 0, 'biassig': 0.027911892, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '1250': {'Kmax': 1.195375, 'lam': -0.032423194, 'sigGs': 5.177596, 'sigGssig': 0.092677385, 'sigTL': 2.708437, 'sigTLsig': 0.07622563, 'sigR': 0.8177013, 'sigRsig': 0.06162229, 'bias': 0, 'biassig': 0.03293372, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '1600': {'Kmax': 1.53008, 'lam': -0.0441045, 'sigGs': 6.29925, 'sigGssig': 0.1153261, 'sigTL': 3.2283993, 'sigTLsig': 0.09118158, 'sigR': 0.988786, 'sigRsig': 0.078567736, 'bias': 0, 'biassig': 0.03877672, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '2000': {'Kmax': 1.9126, 'lam': -0.012963797, 'sigGs': 2.653871, 'sigGssig': 0.015890995, 'sigTL': 1.4356787, 'sigTLsig': 0.02178686, 'sigR': 0.33124214, 'sigRsig': 0.018801652, 'bias': 0, 'biassig': 0.01570677, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '2500': {'Kmax': 2.39075, 'lam': -0.027097283, 'sigGs': 3.200225, 'sigGssig': 0.019307792, 'sigTL': 1.6897862, 'sigTLsig': 0.025873765, 'sigR': 0.38264316, 'sigRsig': 0.023769397, 'bias': 0, 'biassig': 0.018728448, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '3200': {'Kmax': 3.06016, 'lam': -0.034863412, 'sigGs': 3.9193838, 'sigGssig': 0.02649232, 'sigTL': 2.0417721, 'sigTLsig': 0.032873377, 'sigR': 0.44543457, 'sigRsig': 0.030114045, 'bias': 0, 'biassig': 0.021355819, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '4000': {'Kmax': 3.8252, 'lam': -0.043700505, 'sigGs': 4.8015847, 'sigGssig': 0.03781628, 'sigTL': 2.4629273, 'sigTLsig': 0.042401053, 'sigR': 0.52347374, 'sigRsig': 0.03929801, 'bias': 0, 'biassig': 0.026152484, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '5000': {'Kmax': 4.7815, 'lam': -0.053150143, 'sigGs': 5.8995814, 'sigGssig': 0.0625814, 'sigTL': 2.9761007, 'sigTLsig': 0.061326735, 'sigR': 0.6190265, 'sigRsig': 0.05335372, 'bias': 0, 'biassig': 0.058574405, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '6400': {'Kmax': 6.12032, 'lam': -0.07517104, 'sigGs': 7.1163535, 'sigGssig': 0.08435366, 'sigTL': 3.4502964, 'sigTLsig': 0.08226275, 'sigR': 0.7218788, 'sigRsig': 0.0642334, 'bias': 0, 'biassig': 0.059074216, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '8000': {'Kmax': 7.6504, 'lam': -0.08208357, 'sigGs': 8.916516, 'sigGssig': 0.12763213, 'sigTL': 4.269624, 'sigTLsig': 0.13381928, 'sigR': 0.87760293, 'sigRsig': 0.07389065, 'bias': 0, 'biassig': 0.084842026, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '10000': {'Kmax': 9.563, 'lam': -0.073289566, 'sigGs': 11.291476, 'sigGssig': 0.1639773, 'sigTL': 5.495318, 'sigTLsig': 0.16279395, 'sigR': 1.0522343, 'sigRsig': 0.094359785, 'bias': 0, 'biassig': 0.107438326, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '12800': {'Kmax': 12.24064, 'lam': -0.06495205, 'sigGs': 14.245901, 'sigGssig': 0.17283991, 'sigTL': 7.038261, 'sigTLsig': 0.18822834, 'sigR': 1.2749791, 'sigRsig': 0.120479785, 'bias': 0, 'biassig': 0.0944684, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '16000': {'Kmax': 15.3008, 'lam': -0.060692135, 'sigGs': 17.833515, 'sigGssig': 0.19809262, 'sigTL': 8.877547, 'sigTLsig': 0.23338738, 'sigR': 1.5559287, 'sigRsig': 0.15791349, 'bias': 0, 'biassig': 0.09725099, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '20000': {'Kmax': 19.126, 'lam': -0.060213074, 'sigGs': 22.084776, 'sigGssig': 0.21820943, 'sigTL': 11.002351, 'sigTLsig': 0.28806436, 'sigR': 1.8810822, 'sigRsig': 0.18937257, 'bias': 0, 'biassig': 0.4984733, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512}, '25600': {'Kmax': 24.48128, 'lam': -0.09089118, 'sigGs': 25.853043, 'sigGssig': 0.35371417, 'sigTL': 12.175712, 'sigTLsig': 0.4215717, 'sigR': 2.2760193, 'sigRsig': 0.2609267, 'bias': 0, 'biassig': 0.37568903, 'q': 6.103515625e-05, 'wp': 16383, 'bl': 512} } cam_noisy_params['NikonD850'] = { '800':{ 'Kmax':3.320, 'sigGs':4.858579, 'sigGssig':0.002096, 'sigTL':1.509532, 'lam':-0.26, 'sigR':0.263432, 'q':1/(2**14), 'wp':16383, 'bl':512, 'bias':np.array([0,0,0,0]) }, '1600':{ 'Kmax':6.305, 'sigGs':8.695116, 'sigGssig':0.06967, 'sigTL':2.699073, 'lam':-0.26, 'sigR':0.449245, 'q':1/(2**14), 'wp':16383, 'bl':512, 'bias':np.array([0,0,0,0]) }, '3200':{ 'Kmax':11.975, 'sigGs':15.514215, 'sigGssig':0.06967, 'sigTL':4.825994, 'lam':-0.26, 'sigR':0.766122, 'q':1/(2**14), 'wp':16383, 'bl':512, 'bias':np.array([0,0,0,0]) } } cam_noisy_params['IMX686'] = { '100':{ 'Kmax':0.083805, 'sigGs':0.6926457, 'sigGssig':0.002096, 'sigTL':0.67998, 'lam':0.10229621, 'sigR':0.23668, 'q':1/(2**10), 'wp':1023, 'bl':64, 'bias':np.array([0,0,0,0]) }, '6400':{ 'Kmax':8.74253, 'sigGs':12.8901, 'sigGssig':0.06967, 'sigTL':12.8901, 'lam':0.10229621, 'sigR':0, 'q':1/(2**10), 'wp':1023, 'bl':64, 'bias':np.array([-0.08113494,-0.04906388,-0.9408157,-1.2048522]) } } if camera_type in cam_noisy_params: return cam_noisy_params[camera_type][iso] else: # log(f'''Warning: we have not test the noisy parameters of camera "{camera_type}".''') return None # 取最大对应相机的最大ISO生成噪声参数 def sample_params_max(camera_type='NikonD850', ratio=None, iso=None): # 获取已经测算好的相机噪声参数 params = None if iso is not None: params = get_specific_noise_params(camera_type=camera_type, iso=iso) if params is None: if camera_type in Dual_ISO_Cameras: choice = np.random.randint(2) camera_type += '_lowISO' if choice<1 else '_highISO' params = get_camera_noisy_params(camera_type=camera_type) # 根据最小二乘法得到的噪声参数回归模型采样噪声参数 bias = 0 if 'K_ISO' in params and iso is not None: K_ISO = params['K_ISO'] K = np.array(K_ISO[0] * iso + K_ISO[1]) * (1 + np.random.uniform(low=-0.02, high=+0.02)) log_K = np.log(K) else: log_K = params['Kmax'] * (1 + np.random.uniform(low=-0.02, high=+0.02)) # 增加一些扰动，以防测的不准 K = np.exp(log_K) mu_TL = params['sigTLk']*log_K + params['sigTLb'] mu_R = params['sigRk']*log_K + params['sigRb'] mu_Gs = params['sigGsk']*log_K + params['sigGsb'] if 'sigGsk' in params else 2**(-14) # 去掉log sigTL = np.exp(mu_TL) sigR = np.exp(mu_R) sigGs = np.exp(np.random.normal(loc=mu_Gs, scale=params['sigGssig']) if 'sigGssig' in params else mu_Gs) else: K = params['Kmax'] * (1 + np.random.uniform(low=-0.01, high=+0.01)) # 增加一些扰动，以防测的不准 sigGs = np.random.normal(loc=params['sigGs'], scale=params['sigGssig']) if 'sigGssig' in params else params['sigGs'] sigTL = np.random.normal(loc=params['sigTL'], scale=params['sigTLsig']) if 'sigTLsig' in params else params['sigTL'] sigR = np.random.normal(loc=params['sigR'], scale=params['sigRsig']) if 'sigRsig' in params else params['sigR'] bias = params['bias'] wp = params['wp'] bl = params['bl'] lam = params['lam'] q = params['q'] if ratio is None: if 'SonyA7S2' in camera_type: ratio = np.random.uniform(low=100, high=300) else: log_ratio = np.random.uniform(low=0, high=2.08) ratio = np.exp(log_ratio) return {'K':K, 'sigTL':sigTL, 'sigR':sigR, 'sigGs':sigGs, 'bias':bias, 'lam':lam, 'q':q, 'ratio':ratio, 'wp':wp, 'bl':bl} # 噪声参数采样 def sample_params(camera_type='NikonD850', ln_ratio=False): choice = 1 Dual_ISO_Cameras = ['SonyA7S2'] if camera_type in Dual_ISO_Cameras: choice = np.random.randint(2) camera_type += '_lowISO' if choice<1 else '_highISO' # 获取已经测算好的相机噪声参数 params = get_camera_noisy_params(camera_type=camera_type) wp = params['wp'] bl = params['bl'] lam = params['lam'] q = params['q'] Point_ISO_Cameras = ['CRVD', 'BM3D'] if camera_type in Point_ISO_Cameras: if camera_type == 'CRVD': iso_list = [1600,3200,6400,12800,25600] a_list = np.array([3.513262,6.955588,13.486051,26.585953,52.032536]) b_list = np.array([11.917691,38.117816,130.818508,484.539790,1819.818657]) bias_points = np.array([-1.12660, -1.69546, -3.25935, -6.68111, -12.66876]) K_points = np.log(a_list) Gs_points = np.log(np.sqrt(b_list)) choice = np.random.randint(5) # 取点噪声参数 log_K = K_points[choice] K = a_list[choice] mu_TL = params['sigTLk']*log_K + params['sigTLb'] if 'sigTLk' in params else 0 mu_R = params['sigRk']*log_K + params['sigRb'] if 'sigRk' in params else 0 mu_Gs = Gs_points[choice] bias = bias_points[choice] else: # 根据最小二乘法得到的噪声参数回归模型采样噪声参数 bias = 0 log_K = np.random.uniform(low=params['Kmin'], high=params['Kmax']) K = np.exp(log_K) mu_TL = params['sigTLk']*log_K + params['sigTLb'] if 'sigTLk' in params else q mu_R = params['sigRk']*log_K + params['sigRb'] if 'sigRk' in params else q mu_Gs = params['sigGsk']*log_K + params['sigGsb'] if 'sigGsk' in params else q mu_bias = params['uReadk']*log_K + params['uReadb'] log_sigTL = np.random.normal(loc=mu_TL, scale=params['sigTLsig']) if 'sigTLk' in params else 0 log_sigR = np.random.normal(loc=mu_R, scale=params['sigRsig']) if 'sigRk' in params else 0 log_sigGs = np.random.normal(loc=mu_Gs, scale=params['sigGssig']) if 'sigGsk' in params else q log_bias = np.random.normal(loc=mu_bias, scale=params['uReadsig']) if 'uReadk' in params else 0 # 去掉log sigTL = np.exp(log_sigTL) sigR = np.exp(log_sigR) sigGs = np.exp(log_sigGs) bias = np.exp(log_bias) # 模拟曝光衰减的系数, ln_ratio模式会照顾弱噪声场景，更有通用性 if ln_ratio: high = 1 if 'CRVD' in camera_type else 5 log_ratio = np.random.uniform(low=-0.01, high=high) ratio = np.exp(log_ratio) # np.random.uniform(low=1, high=200) if choice else np.exp(log_ratio) else: ratio = np.random.uniform(low=100, high=300) return {'K':K, 'sigTL':sigTL, 'sigR':sigR, 'sigGs':sigGs, 'bias':bias, 'lam':lam, 'q':q, 'ratio':ratio, 'wp':wp, 'bl':bl} def get_aug_param_torch(data, b=8, command='augv5', numpy=False, camera_type='SonyA7S2'): aug_r, aug_g, aug_b = torch.zeros(b), torch.zeros(b), torch.zeros(b) r = np.random.randint(2) * 0.25 + 0.25 u = r if np.random.randint(4): if 'augv5' in command: # v5，依照相机白平衡经验参数增强，激进派作风 rgb_gain, red_gain, blue_gain = random_gains(camera_type) rgb_gain = 1 / rgb_gain rg = data["wb"][:, 0] / red_gain[0] bg = data["wb"][:, 2] / blue_gain[0] aug_g = torch.rand(b) * r + rgb_gain[0] - 0.9 aug_r = torch.rand(b) * r + rg * (1+aug_g) - 1.1 aug_b = torch.rand(b) * r + bg * (1+aug_g) - 1.1 elif 'augv2' in command: # v2，相信原数据颜色的准确性，保守派作风 aug_g = torch.clamp(torch.randn(b) * r, 0, 4*u) aug_r = torch.clamp((1+torch.randn(b)*r) * (1+aug_g) - 1, 0, 4*u) aug_b = torch.clamp((1+torch.randn(b)*r) * (1+aug_g) - 1, 0, 4*u) # 保证非负 daug, _ = torch.min(torch.stack((aug_r, aug_g, aug_b)), dim=0) daug[daug>0] = 0 aug_r = (1+aug_r) / (1+daug) - 1 aug_g = (1+aug_g) / (1+daug) - 1 aug_b = (1+aug_b) / (1+daug) - 1 if numpy: aug_r = np.squeeze(aug_r.numpy()) aug_g = np.squeeze(aug_g.numpy()) aug_b = np.squeeze(aug_b.numpy()) return aug_r, aug_g, aug_b def raw_wb_aug(noisy, gt, aug_wb=None, camera_type='SonyA7S2', ratio=1, ori=True, iso=None): # [c, h, w] p = get_specific_noise_params(camera_type=camera_type, iso=iso) if p is None: assert camera_type == 'SonyA7S2' camera_type += '_lowISO' if iso<=1600 else '_highISO' p = get_camera_noisy_params(camera_type=camera_type) # 根据最小二乘法得到的噪声参数回归模型采样噪声参数 p['K'] = 0.0009546 * iso * (1 + np.random.uniform(low=-0.01, high=+0.01)) - 0.00193 log_K = np.log(p['K']) mu_Gs = p['sigGsk']*log_K + p['sigGsb'] p['sigGs'] = np.exp(np.random.normal(loc=mu_Gs, scale=p['sigGssig'])) else: p['K'] = p['Kmax'] * (1 + np.random.uniform(low=-0.01, high=+0.01)) # 增加一些扰动，以防测的不准 p['sigGs'] = np.random.normal(loc=p['sigGs'], scale=p['sigGssig']) if 'sigGssig' in p else p['sigGs'] if aug_wb is not None: # 默认pattern为RGGB！(通道rgbg排列) gt = gt * (p['wp'] - p['bl']) / ratio noisy = noisy * (p['wp'] - p['bl']) # 补噪声 daug = -np.minimum(np.min(aug_wb), 0) # aug_wb = aug_wb + daug if daug == 0: # 只有增益的话很好处理，叠加泊松分布就行 dy = gt * aug_wb.reshape(-1,1,1) # 我考虑过这里dy和dn要不要量化一下，感觉不量化对多样性更友好 dn = np.random.poisson(dy/p['K']).astype(np.float32) * p['K'] else: # BiSNA，折中方案，不好讲故事，弃疗了 raise NotImplementedError # 存在减益的话就很麻烦，需要考虑read noise并且补齐分布 scale = 1 - daug # 要保证dyn是非负的 aug_wb_new = aug_wb + daug dy = gt * aug_wb.reshape(-1,1,1) dyn = gt * aug_wb_new.reshape(-1,1,1) # 先通过缩放减小噪声图 noisy *= scale # 补齐单个照片的读噪声 dn_read = np.random.randn(*gt.shape).astype(np.float32) * p['sigGs'] * np.sqrt(1-scale**2) # 补齐由于除法导致的分布变化，恢复shot noise应有的分布（不完全相等） scale_sigma = scale - scale**2 dn_shot = np.random.poisson(scale_sigma * gt/p['K']).astype(np.float32)*p['K'] - gt * scale_sigma # 叠加泊松分布 dn_aug = np.random.poisson(dyn/p['K']).astype(np.float32) * p['K'] dn = dn_read + dn_shot + dn_aug gt = np.clip((gt + dy)*ratio, 0, (p['wp'] - p['bl'])) noisy = np.clip(noisy + dn, -p['bl'], (p['wp'] - p['bl'])) gt /= (p['wp'] - p['bl']) noisy /= (p['wp'] - p['bl']) if ori is False: noisy *= ratio return noisy.astype(np.float32), gt.astype(np.float32) def raw_wb_aug_torch(noisy, gt, aug_wb=None, camera_type='IMX686', ratio=1, ori=True, iso=None, ratiofix=False): p = get_specific_noise_params(camera_type=camera_type, iso=iso) if p is None: assert camera_type == 'SonyA7S2' camera_type += '_lowISO' if iso<=1600 else '_highISO' p = get_camera_noisy_params(camera_type=camera_type) # 根据最小二乘法得到的噪声参数回归模型采样噪声参数 p['K'] = 0.0009546 * iso * (1 + np.random.uniform(low=-0.01, high=+0.01)) - 0.00193 log_K = np.log(p['K']) mu_Gs = p['sigGsk']*log_K + p['sigGsb'] p['sigGs'] = np.exp(np.random.normal(loc=mu_Gs, scale=p['sigGssig'])) else: p['K'] = p['Kmax'] * (1 + np.random.uniform(low=-0.01, high=+0.01)) # 增加一些扰动，以防测的不准 p['sigGs'] = np.random.normal(loc=p['sigGs'], scale=p['sigGssig']) if 'sigGssig' in p else p['sigGs'] if aug_wb is not None: # 默认pattern为RGGB！(通道rgbg排列) gt = gt * (p['wp'] - p['bl']) / ratio noisy = noisy * (p['wp'] - p['bl']) # 补噪声 daug = -np.minimum(np.min(aug_wb), 0) daug = torch.from_numpy(np.array(daug)).to(noisy.device) aug_wb = torch.from_numpy(aug_wb).to(noisy.device) dy = gt * aug_wb.reshape(-1,1,1) # 我考虑过这里dy和dn要不要量化一下，感觉不量化对多样性更友好 if daug == 0: # 只有增益的话很好处理，叠加泊松分布就行 dn = tdist.Poisson(dy/p['K']).sample() * p['K'] else: # BiSNA，折中方案，不好讲故事，弃疗了 warnings.warn('You are using BiSNA!!!') raise NotImplementedError # 存在减益的话就很麻烦，需要考虑read noise并且补齐分布 scale = 1 - daug # 要保证dyn是非负的 aug_wb_new = aug_wb + daug dyn = gt * aug_wb_new.reshape(-1,1,1) # 先通过缩放减小噪声图 noisy *= scale # 补齐单个照片的读噪声 dn_read = tdist.Normal(0, p['sigGs']).sample() * torch.sqrt(1-scale**2) # 补齐由于除法导致的分布变化，恢复shot noise应有的分布（不完全相等） scale_sigma = scale - scale**2 dn_shot = tdist.Poisson(scale_sigma * gt/p['K']).sample() *p['K'] - gt * scale_sigma # 叠加泊松分布 dn_aug = tdist.Poisson(dyn/p['K']).sample() * p['K'] dn = dn_read + dn_shot + dn_aug if ratiofix: ratio = ratio / (1 + daug) gt = torch.clamp((gt + dy)*ratio, 0, (p['wp'] - p['bl'])) noisy = torch.clamp(noisy + dn, -p['bl'], (p['wp'] - p['bl'])) gt /= (p['wp'] - p['bl']) noisy /= (p['wp'] - p['bl']) if ori is False: noisy *= ratio return noisy, gt def SNA_torch(gt, aug_wb, camera_type='IMX686', ratio=1, black_lr=False, ori=True, iso=None): p = get_specific_noise_params(camera_type=camera_type, iso=iso) if p is None: assert camera_type == 'SonyA7S2' camera_type += '_lowISO' if iso<=1600 else '_highISO' p = get_camera_noisy_params(camera_type=camera_type) # 根据最小二乘法得到的噪声参数回归模型采样噪声参数 p['K'] = 0.0009546 * iso * (1 + np.random.uniform(low=-0.01, high=+0.01)) - 0.00193 else: p['K'] = p['Kmax'] * (1 + np.random.uniform(low=-0.01, high=+0.01)) # 增加一些扰动，以防测的不准 # 默认pattern为RGGB！(通道rgbg排列) gt = gt * (p['wp'] - p['bl']) / ratio # 补噪声 aug_wb = torch.from_numpy(aug_wb).to(gt.device) dy = gt * aug_wb.reshape(-1,1,1) # 我考虑过这里dy和dn要不要量化一下，感觉不量化对多样性更友好 # 只有增益的话很好处理，叠加泊松分布就行 dn = tdist.Poisson(dy/p['K']).sample() * p['K'] # 贴黑图的，所以抛去gt中多算的一份泊松分布 if black_lr: dy = dy - gt dy = dy * ratio / (p['wp'] - p['bl']) dn = dn / (p['wp'] - p['bl']) if ori is False: dn *= ratio return dn, dy # @ fn_timer def generate_noisy_obs(y, camera_type=None, noise_code='p', param=None, MultiFrameMean=1, ori=False, clip=False): p = param y = y * (p['wp'] - p['bl']) # p['ratio'] = 1/p['ratio'] # 临时行为，为了快速实现MFM y = y / p['ratio'] MFM = MultiFrameMean ** 0.5 use_R = True if 'r' in noise_code.lower() else False use_Q = True if 'q' in noise_code.lower() else False use_TL = True if 'g' in noise_code.lower() else False use_P = True if 'p' in noise_code.lower() else False use_D = True if 'd' in noise_code.lower() else False use_black = True if 'b' in noise_code.lower() else False if use_P: # 使用泊松噪声作为shot noisy noisy_shot = np.random.poisson(MFM*y/p['K']).astype(np.float32) * p['K'] / MFM else: # 不考虑shot noisy noisy_shot = y + np.random.randn(*y.shape).astype(np.float32) * np.sqrt(np.maximum(y/p['K'], 1e-10)) * p['K'] / MFM if not use_black: if use_TL: # 使用TL噪声作为read noisy noisy_read = stats.tukeylambda.rvs(p['lam'], scale=p['sigTL']/MFM, size=y.shape).astype(np.float32) else: # 使用高斯噪声作为read noisy noisy_read = stats.norm.rvs(scale=p['sigGs']/MFM, size=y.shape).astype(np.float32) # 行噪声需要使用行的维度h，[1,c,h,w]所以-2是h noisy_row = np.random.randn(y.shape[-3], y.shape[-2], 1).astype(np.float32) * p['sigR']/MFM if use_R else 0 noisy_q = np.random.uniform(low=-0.5, high=0.5, size=y.shape) if use_Q else 0 noisy_bias = p['bias'].reshape(-1,1,1) if use_D else 0 else: noisy_read = 0 noisy_row = 0 noisy_q = 0 noisy_bias = 0 # 归一化回[0, 1] z = (noisy_shot + noisy_read + noisy_row + noisy_q + noisy_bias) / (p['wp'] - p['bl']) # 模拟实际情况 z = np.clip(z, -p['bl']/p['wp'], 1) if not clip else np.clip(z, 0, 1) if ori is False: z = z * p['ratio'] return z.astype(np.float32) # @fn_timer def generate_noisy_torch(y, camera_type=None, noise_code='p', param=None, MultiFrameMean=1, ori=False, clip=False): p = param y = y * (p['wp'] - p['bl']) # p['ratio'] = 1/p['ratio'] # 临时行为，为了快速实现MFM y = y / p['ratio'] MFM = MultiFrameMean ** 0.5 use_R = True if 'r' in noise_code.lower() else False use_Q = True if 'q' in noise_code.lower() else False use_TL = True if 'g' in noise_code.lower() else False use_P = True if 'p' in noise_code.lower() else False use_D = True if 'd' in noise_code.lower() else False use_black = True if 'b' in noise_code.lower() else False if use_P: # 使用泊松噪声作为shot noisy noisy_shot = tdist.Poisson(MFM*y/p['K']).sample() * p['K'] / MFM else: # 不考虑shot noisy noisy_shot = tdist.Normal(y).sample() * torch.sqrt(torch.max(y/p['K'], 1e-10)) * p['K'] / MFM if not use_black: if use_TL: # 使用TL噪声作为read noisy raise NotImplementedError # noisy_read = stats.tukeylambda.rvs(p['lam'], scale=p['sigTL'], size=y.shape).astype(np.float32) else: # 使用高斯噪声作为read noisy noisy_read = tdist.Normal(loc=torch.zeros_like(y), scale=p['sigGs']/MFM).sample() else: noisy_read = 0 # 使用行噪声 noisy_row = torch.randn(y.shape[-3], y.shape[-2], 1, device=y.device) * p['sigR'] / MFM if use_R else 0 noisy_q = (torch.rand(y.shape, device=y.device) - 0.5) * p['q'] * (p['wp'] - p['bl']) if use_Q else 0 noisy_bias = torch.from_numpy(p['bias'].reshape(-1,1,1)) if use_D else 0 # 归一化回[0, 1] z = (noisy_shot + noisy_read + noisy_row + noisy_q + noisy_bias) / (p['wp'] - p['bl']) # 模拟实际情况 z = torch.clamp(z, -p['bl']/p['wp'], 1) if not clip else torch.clamp(z, 0, 1) # ori_brightness if ori is False: z = z * p['ratio'] return z class HighBitRecovery: def __init__(self, camera_type='IMX686', noise_code='prq', param=None, perturb=True, factor=6, float=True): self.camera_type = camera_type self.noise_code = noise_code self.param = param self.perturb = perturb self.factor = factor self.float = float self.lut = {} def get_lut(self, iso_list, blc_mean=None): for iso in iso_list: if blc_mean is None: bias = 0 else: bias = np.mean(blc_mean[iso]) if self.perturb: sigma_t = 0.1 bias += np.random.randn() * sigma_t self.lut[iso] = self.HB2LB_LUT(iso, bias) def HB2LB_LUT(self, iso, bias=0, param=None): # 标记LUT区间 lut_info = {} # 获得标定的噪声参数 p = sample_params_max(self.camera_type, iso=iso) if param is None else param lut_info['param'] = p # 选择一种分布，依照该分布映射到HighBit if 'g' in self.noise_code.lower(): dist = stats.tukeylambda(p['lam'], loc=bias, scale=p['sigTL']) sigma = p['sigTL'] lut_info['dist'] = dist else: dist = stats.norm(loc=bias, scale=p['sigGs']) sigma = p['sigGs'] lut_info['dist'] = dist # 寻址范围为[-6σ,6σ]，离群点不做映射恢复 low = max(int(-sigma*self.factor + bias), -p['bl']+1) high = int(sigma*self.factor + bias) for x in range(low, high): lut_info[x] = { # 累积概率函数的起点 'cdf': dist.cdf(x-0.5), # 累积概率函数的变化范围 'range': dist.cdf(x+0.5) - dist.cdf(x-0.5), } lut_info['low'] = low lut_info['high'] = high lut_info['bias'] = bias lut_info['sigma'] = sigma return lut_info def map(self, data, iso=6400, norm=True): # 将LB图像映射成HB图像 p = self.lut[iso]['param'] if np.max(data) <= 1: data = data * (p['wp'] - p['bl']) data_float = data.copy() data = np.round(data_float) if self.float: delta = data_float - data rand = np.random.uniform(0, 1, size=data.shape) # 快速版：寻址范围为[-6σ,6σ]，离群点不做映射恢复 for x in range(self.lut[iso]['low'], self.lut[iso]['high']): keys = (data==x) cdf = self.lut[iso][x]['cdf'] r = self.lut[iso][x]['range'] # 根据ppf反推分位点 data[keys] = self.lut[iso]['dist'].ppf(cdf + rand[keys] * r) if self.float: data = data + delta if norm: data = data / (p['wp'] - p['bl']) else: data = data + p['bl'] return data class ELDEvalDataset(torch.utils.data.Dataset): # 初始化函数，传入参数：basedir（基础路径），camera_suffix（相机后缀），scenes（场景），img_ids（图像id） def __init__(self, basedir, camera_suffix=('NikonD850','.nef'), scenes=None, img_ids=None): super(ELDEvalDataset, self).__init__() self.basedir = basedir self.camera_suffix = camera_suffix # ('Canon', '.CR2') self.scenes = scenes self.img_ids = img_ids # self.input_dict = {} # self.target_dict = {} # 获取每一个图像的数据 def __getitem__(self, i): camera, suffix = self.camera_suffix scene_id = i // len(self.img_ids) img_id = i % len(self.img_ids) scene ='scene-{}'.format(self.scenes[scene_id]) datadir = os.path.join(self.basedir, camera, scene) input_path = os.path.join(datadir, 'IMG_{:04d}{}'.format(self.img_ids[img_id], suffix)) gt_ids = np.array([1, 6, 11, 16]) ind = np.argmin(np.abs(self.img_ids[img_id] - gt_ids)) target_path = os.path.join(datadir, 'IMG_{:04d}{}'.format(gt_ids[ind], suffix)) iso, expo = metainfo(target_path) target_expo = iso * expo iso, expo = metainfo(input_path) ratio = target_expo / (iso * expo) with rawpy.imread(input_path) as raw: input = pack_raw_bayer(raw) * ratio with rawpy.imread(target_path) as raw: target = pack_raw_bayer(raw) input = np.maximum(np.minimum(input, 1.0), 0) target = np.maximum(np.minimum(target, 1.0), 0) input = np.ascontiguousarray(input) target = np.ascontiguousarray(target) data = {'input': input, 'target': target, 'fn':input_path, 'rawpath': target_path} return data # 获取数据的长度 def __len__(self): return len(self.scenes) * len(self.img_ids) if __name__ == '__main__': path = 'F:/datasets/ELD/SonyA7S2/scene-8' files = [os.path.join(path, name) for name in os.listdir(path) if '.ARW' in name] for name in files: print(name) raw = rawpy.imread(name) img = raw.raw_image_visible.astype(np.float32)[np.newaxis,:,:] # img = img[:, 1000:1500, 2200:2700] fig = plt.figure(figsize=(16,10)) img = np.clip((img-512) / (16383-512), 0, 1) p = sample_params_max(camera_type='SonyA7S2', ratio=100, iso=1600) noisy = generate_noisy_obs(img,camera_type='SonyA7S2', param=p) # refer_path = path+'\\'+'DSC02750.ARW' # raw_refer = rawpy.imread(refer_path) # print(np.min(raw_refer.raw_image_visible), np.max(raw_refer.raw_image_visible), np.mean(raw_refer.raw_image_visible)) # raw_refer.raw_image_visible[:,:] = np.clip((raw_refer.raw_image_visible.astype(np.float32)-512) / (16383-512)*200, 0, 1)*16383 # print(np.min(raw_refer.raw_image_visible), np.max(raw_refer.raw_image_visible), np.mean(raw_refer.raw_image_visible)) # out1 = raw_refer.postprocess(use_camera_wb=True, no_auto_bright=True) # print(np.min(out1), np.max(out1), np.mean(out1)) # plt.imsave('real.png', out1) # plt.imshow(out1) # plt.show() raw.raw_image_visible[:,:] = noisy[0,:,:]*16383 out = raw.postprocess(use_camera_wb=True) plt.imshow(out) plt.show() plt.imsave('gen.png', out) print('test') print("")