borealisai / noise_flow Goto Github PK

from borealisflows.NoiseFlowWrapper import NoiseFlowWrapper

patch_size, stride = 32, 32  # patch size  = [32, 32, 4]
aug_times = 1
scales = [1]  # [1, 0.9, 0.8, 0.7]
nf_model_path = 'models/NoiseFlow'

def load_cam_iso_nlf():
    cin = pd.read_csv('cam_iso_nlf.txt')
    cin = cin.drop_duplicates()
    cin = cin.set_index('cam_iso', drop=False)
    return cin

# Prepare NoiseFlow
noise_flow = NoiseFlowWrapper(nf_model_path)

# camera IDs and ISO levels related to the SIDD dataset
cam_iso_nlf = load_cam_iso_nlf()
n_cam_iso = cam_iso_nlf['cam_iso'].count()
iso_vals = [100.0, 400.0, 800.0, 1600.0, 3200.0]
cam_ids = [0, 1, 3, 3, 4]  # IP, GP, S6, N6, G4
cam_vals = ['IP', 'GP', 'S6', 'N6', 'G4']


def pack_raw(rgb_img):
    """Packs sRGB image to 4 channels (h, w) --> (h/2, w/2, 4)."""
    # pack sRGB image to 4 channels
    im = np.expand_dims(rgb_img, axis=2)
    img_shape = im.shape
    h = img_shape[0]
    w = img_shape[1]
    out = np.concatenate((im[0:h, 0:w, :, 0],
                          im[0:h, 0:w, :, 1],
                          im[0:h, 0:w, :, 1],
                          im[0:h, 0:w, :, 2]), axis=2)
    return out

def gen_patches_png(img_file):
    # split patches of a given img
    img = img_file
    img = np.expand_dims(pack_raw(img), axis=0)
    _,h, w, c = img.shape
    patches = None
    # extract patches
    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            x = img[0,i:i + patch_size, j:j + patch_size, :]  # first dim will be removed
            for k in range(0, aug_times):
                if patches is None:
                    patches = x[np.newaxis, :, :, :]  # restore first dim
                else:
                    patches = np.concatenate((patches, x[np.newaxis, :, :, :]), axis=0) 
    return patches,[h,w]

def recover_full_img(patches,shape):
   #Combines patches together into an img
    h,w = shape
    out_image = np.zeros((h,w,3))
    index=0
    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            out_image[i:i + patch_size,j:j + patch_size,0] = patches[index,:,:,0]
            out_image[i:i + patch_size,j:j + patch_size,1] = patches[index,:,:,1]
            out_image[i:i + patch_size,j:j + patch_size,2] = patches[index,:,:,3]
            index = index+1
    return out_image


img_clean = imageio.imread(img_file).astype('float32')/255.0
patches,img_shape = gen_patches_png(img_clean)
patches = patches**2.2  ##sRGB to linear
cam_iso_idx = random.choice([3,9]) ## ISO-1600
row = cam_iso_nlf.iloc[cam_iso_idx]
cam = cam_vals.index(row['cam_iso'][:2])
iso = float(row['cam_iso'][3:])
noise = noise_flow.sample_noise_nf(patches, 0.0, 0.0, iso, cam)
patches_noisy = np.clip(patches + noise,0,1)
patches_noisy = patches_noisy**(1/2.2) ##linear to sRGB
recovered_img_noisy = recover_full_img(patches_noisy,img_shape)
imageio.imwrite(cur_img.replace('.png','_NF.png'),recovered_img_noisy)