Spaces:
Running
Running
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.autograd import Variable | |
| import torch | |
| import numpy as np | |
| import os, time, random | |
| import argparse | |
| from torch.utils.data import Dataset, DataLoader | |
| from PIL import Image as PILImage | |
| from glob import glob | |
| from tqdm import tqdm | |
| import rawpy | |
| import colour_demosaicing | |
| from .InvISP.model.model import InvISPNet | |
| from .utils.common import Notify | |
| from datasets.noise import ( | |
| camera_params, | |
| addGStarNoise, | |
| addPStarNoise, | |
| addQuantNoise, | |
| addRowNoise, | |
| sampleK, | |
| ) | |
| class NoiseSimulator: | |
| def __init__(self, device, ckpt_path="./datasets/InvISP/pretrained/canon.pth"): | |
| self.device = device | |
| # load Invertible ISP Network | |
| self.net = ( | |
| InvISPNet(channel_in=3, channel_out=3, block_num=8).to(self.device).eval() | |
| ) | |
| self.net.load_state_dict(torch.load(ckpt_path), strict=False) | |
| print( | |
| Notify.INFO, "Loaded ISPNet checkpoint: {}".format(ckpt_path), Notify.ENDC | |
| ) | |
| # white balance parameters | |
| self.wb = np.array([2020.0, 1024.0, 1458.0, 1024.0]) | |
| # use Canon EOS 5D4 noise parameters provided by ELD | |
| self.camera_params = camera_params | |
| # random specify exposure time ratio from 50 to 150 | |
| self.ratio_min = 50 | |
| self.ratio_max = 150 | |
| pass | |
| # inverse demosaic | |
| # input: [H, W, 3] | |
| # output: [H, W] | |
| def invDemosaic(self, img): | |
| img_R = img[::2, ::2, 0] | |
| img_G1 = img[::2, 1::2, 1] | |
| img_G2 = img[1::2, ::2, 1] | |
| img_B = img[1::2, 1::2, 2] | |
| raw_img = np.ones(img.shape[:2]) | |
| raw_img[::2, ::2] = img_R | |
| raw_img[::2, 1::2] = img_G1 | |
| raw_img[1::2, ::2] = img_G2 | |
| raw_img[1::2, 1::2] = img_B | |
| return raw_img | |
| # demosaic - nearest ver | |
| # input: [H, W] | |
| # output: [H, W, 3] | |
| def demosaicNearest(self, img): | |
| raw = np.ones((img.shape[0], img.shape[1], 3)) | |
| raw[::2, ::2, 0] = img[::2, ::2] | |
| raw[::2, 1::2, 0] = img[::2, ::2] | |
| raw[1::2, ::2, 0] = img[::2, ::2] | |
| raw[1::2, 1::2, 0] = img[::2, ::2] | |
| raw[::2, ::2, 2] = img[1::2, 1::2] | |
| raw[::2, 1::2, 2] = img[1::2, 1::2] | |
| raw[1::2, ::2, 2] = img[1::2, 1::2] | |
| raw[1::2, 1::2, 2] = img[1::2, 1::2] | |
| raw[::2, ::2, 1] = img[::2, 1::2] | |
| raw[::2, 1::2, 1] = img[::2, 1::2] | |
| raw[1::2, ::2, 1] = img[1::2, ::2] | |
| raw[1::2, 1::2, 1] = img[1::2, ::2] | |
| return raw | |
| # demosaic | |
| # input: [H, W] | |
| # output: [H, W, 3] | |
| def demosaic(self, img): | |
| return colour_demosaicing.demosaicing_CFA_Bayer_bilinear(img, "RGGB") | |
| # load rgb image | |
| def path2rgb(self, path): | |
| return torch.from_numpy(np.array(PILImage.open(path)) / 255.0) | |
| # InvISP | |
| # input: rgb image [H, W, 3] | |
| # output: raw image [H, W] | |
| def rgb2raw(self, rgb, batched=False): | |
| # 1. rgb -> invnet | |
| if not batched: | |
| rgb = rgb.unsqueeze(0) | |
| rgb = rgb.permute(0, 3, 1, 2).float().to(self.device) | |
| with torch.no_grad(): | |
| reconstruct_raw = self.net(rgb, rev=True) | |
| pred_raw = reconstruct_raw.detach().permute(0, 2, 3, 1) | |
| pred_raw = torch.clamp(pred_raw, 0, 1) | |
| if not batched: | |
| pred_raw = pred_raw[0, ...] | |
| pred_raw = pred_raw.cpu().numpy() | |
| # 2. -> inv gamma | |
| norm_value = np.power(16383, 1 / 2.2) | |
| pred_raw *= norm_value | |
| pred_raw = np.power(pred_raw, 2.2) | |
| # 3. -> inv white balance | |
| wb = self.wb / self.wb.max() | |
| pred_raw = pred_raw / wb[:-1] | |
| # 4. -> add black level | |
| pred_raw += self.camera_params["black_level"] | |
| # 5. -> inv demosaic | |
| if not batched: | |
| pred_raw = self.invDemosaic(pred_raw) | |
| else: | |
| preds = [] | |
| for i in range(pred_raw.shape[0]): | |
| preds.append(self.invDemosaic(pred_raw[i])) | |
| pred_raw = np.stack(preds, axis=0) | |
| return pred_raw | |
| def raw2noisyRaw(self, raw, ratio_dec=1, batched=False): | |
| if not batched: | |
| ratio = (random.uniform(self.ratio_min, self.ratio_max) - 1) * ratio_dec + 1 | |
| raw = raw.copy() / ratio | |
| K = sampleK(self.camera_params["Kmin"], self.camera_params["Kmax"]) | |
| q = 1 / ( | |
| self.camera_params["max_value"] - self.camera_params["black_level"] | |
| ) | |
| raw = addPStarNoise(raw, K) | |
| raw = addGStarNoise( | |
| raw, | |
| K, | |
| self.camera_params["G_shape"], | |
| self.camera_params["Profile-1"]["G_scale"], | |
| ) | |
| raw = addRowNoise(raw, K, self.camera_params["Profile-1"]["R_scale"]) | |
| raw = addQuantNoise(raw, q) | |
| raw *= ratio | |
| return raw | |
| else: | |
| raw = raw.copy() | |
| for i in range(raw.shape[0]): | |
| ratio = random.uniform(self.ratio_min, self.ratio_max) | |
| raw[i] /= ratio | |
| K = sampleK(self.camera_params["Kmin"], self.camera_params["Kmax"]) | |
| q = 1 / ( | |
| self.camera_params["max_value"] - self.camera_params["black_level"] | |
| ) | |
| raw[i] = addPStarNoise(raw[i], K) | |
| raw[i] = addGStarNoise( | |
| raw[i], | |
| K, | |
| self.camera_params["G_shape"], | |
| self.camera_params["Profile-1"]["G_scale"], | |
| ) | |
| raw[i] = addRowNoise( | |
| raw[i], K, self.camera_params["Profile-1"]["R_scale"] | |
| ) | |
| raw[i] = addQuantNoise(raw[i], q) | |
| raw[i] *= ratio | |
| return raw | |
| def raw2rgb(self, raw, batched=False): | |
| # 1. -> demosaic | |
| if not batched: | |
| raw = self.demosaic(raw) | |
| else: | |
| raws = [] | |
| for i in range(raw.shape[0]): | |
| raws.append(self.demosaic(raw[i])) | |
| raw = np.stack(raws, axis=0) | |
| # 2. -> substract black level | |
| raw -= self.camera_params["black_level"] | |
| raw = np.clip( | |
| raw, 0, self.camera_params["max_value"] - self.camera_params["black_level"] | |
| ) | |
| # 3. -> white balance | |
| wb = self.wb / self.wb.max() | |
| raw = raw * wb[:-1] | |
| # 4. -> gamma | |
| norm_value = np.power(16383, 1 / 2.2) | |
| raw = np.power(raw, 1 / 2.2) | |
| raw /= norm_value | |
| # 5. -> ispnet | |
| if not batched: | |
| input_raw_img = ( | |
| torch.Tensor(raw) | |
| .permute(2, 0, 1) | |
| .float() | |
| .to(self.device)[np.newaxis, ...] | |
| ) | |
| else: | |
| input_raw_img = ( | |
| torch.Tensor(raw).permute(0, 3, 1, 2).float().to(self.device) | |
| ) | |
| with torch.no_grad(): | |
| reconstruct_rgb = self.net(input_raw_img) | |
| reconstruct_rgb = torch.clamp(reconstruct_rgb, 0, 1) | |
| pred_rgb = reconstruct_rgb.detach().permute(0, 2, 3, 1) | |
| if not batched: | |
| pred_rgb = pred_rgb[0, ...] | |
| pred_rgb = pred_rgb.cpu().numpy() | |
| return pred_rgb | |
| def raw2packedRaw(self, raw, batched=False): | |
| # 1. -> substract black level | |
| raw -= self.camera_params["black_level"] | |
| raw = np.clip( | |
| raw, 0, self.camera_params["max_value"] - self.camera_params["black_level"] | |
| ) | |
| raw /= self.camera_params["max_value"] | |
| # 2. pack | |
| if not batched: | |
| im = np.expand_dims(raw, axis=2) | |
| img_shape = im.shape | |
| H = img_shape[0] | |
| W = img_shape[1] | |
| out = np.concatenate( | |
| ( | |
| im[0:H:2, 0:W:2, :], | |
| im[0:H:2, 1:W:2, :], | |
| im[1:H:2, 1:W:2, :], | |
| im[1:H:2, 0:W:2, :], | |
| ), | |
| axis=2, | |
| ) | |
| else: | |
| im = np.expand_dims(raw, axis=3) | |
| img_shape = im.shape | |
| H = img_shape[1] | |
| W = img_shape[2] | |
| out = np.concatenate( | |
| ( | |
| im[:, 0:H:2, 0:W:2, :], | |
| im[:, 0:H:2, 1:W:2, :], | |
| im[:, 1:H:2, 1:W:2, :], | |
| im[:, 1:H:2, 0:W:2, :], | |
| ), | |
| axis=3, | |
| ) | |
| return out | |
| def raw2demosaicRaw(self, raw, batched=False): | |
| # 1. -> demosaic | |
| if not batched: | |
| raw = self.demosaic(raw) | |
| else: | |
| raws = [] | |
| for i in range(raw.shape[0]): | |
| raws.append(self.demosaic(raw[i])) | |
| raw = np.stack(raws, axis=0) | |
| # 2. -> substract black level | |
| raw -= self.camera_params["black_level"] | |
| raw = np.clip( | |
| raw, 0, self.camera_params["max_value"] - self.camera_params["black_level"] | |
| ) | |
| raw /= self.camera_params["max_value"] | |
| return raw | |