Spaces:
Running
Running
| from __future__ import print_function, division | |
| import os, random, time | |
| import torch | |
| import numpy as np | |
| from torch.utils.data import Dataset | |
| from torchvision import transforms, utils | |
| import rawpy | |
| from glob import glob | |
| from PIL import Image as PILImage | |
| import numbers | |
| from scipy.misc import imread | |
| from .base_dataset import BaseDataset | |
| class FiveKDatasetTrain(BaseDataset): | |
| def __init__(self, opt): | |
| super().__init__(opt=opt) | |
| self.patch_size = 256 | |
| input_RAWs_WBs, target_RGBs = self.load(is_train=True) | |
| assert len(input_RAWs_WBs) == len(target_RGBs) | |
| self.data = {'input_RAWs_WBs':input_RAWs_WBs, 'target_RGBs':target_RGBs} | |
| def random_flip(self, input_raw, target_rgb): | |
| idx = np.random.randint(2) | |
| input_raw = np.flip(input_raw,axis=idx).copy() | |
| target_rgb = np.flip(target_rgb,axis=idx).copy() | |
| return input_raw, target_rgb | |
| def random_rotate(self, input_raw, target_rgb): | |
| idx = np.random.randint(4) | |
| input_raw = np.rot90(input_raw,k=idx) | |
| target_rgb = np.rot90(target_rgb,k=idx) | |
| return input_raw, target_rgb | |
| def random_crop(self, patch_size, input_raw, target_rgb,flow=False,demos=False): | |
| H, W, _ = input_raw.shape | |
| rnd_h = random.randint(0, max(0, H - patch_size)) | |
| rnd_w = random.randint(0, max(0, W - patch_size)) | |
| patch_input_raw = input_raw[rnd_h:rnd_h + patch_size, rnd_w:rnd_w + patch_size, :] | |
| if flow or demos: | |
| patch_target_rgb = target_rgb[rnd_h:rnd_h + patch_size, rnd_w:rnd_w + patch_size, :] | |
| else: | |
| patch_target_rgb = target_rgb[rnd_h*2:rnd_h*2 + patch_size*2, rnd_w*2:rnd_w*2 + patch_size*2, :] | |
| return patch_input_raw, patch_target_rgb | |
| def aug(self, patch_size, input_raw, target_rgb, flow=False, demos=False): | |
| input_raw, target_rgb = self.random_crop(patch_size, input_raw,target_rgb,flow=flow, demos=demos) | |
| input_raw, target_rgb = self.random_rotate(input_raw,target_rgb) | |
| input_raw, target_rgb = self.random_flip(input_raw,target_rgb) | |
| return input_raw, target_rgb | |
| def __len__(self): | |
| return len(self.data['input_RAWs_WBs']) | |
| def __getitem__(self, idx): | |
| input_raw_wb_path = self.data['input_RAWs_WBs'][idx] | |
| target_rgb_path = self.data['target_RGBs'][idx] | |
| target_rgb_img = imread(target_rgb_path) | |
| input_raw_wb = np.load(input_raw_wb_path) | |
| input_raw_img = input_raw_wb['raw'] | |
| wb = input_raw_wb['wb'] | |
| wb = wb / wb.max() | |
| input_raw_img = input_raw_img * wb[:-1] | |
| self.patch_size = 256 | |
| input_raw_img, target_rgb_img = self.aug(self.patch_size, input_raw_img, target_rgb_img, flow=True, demos=True) | |
| if self.gamma: | |
| norm_value = np.power(4095, 1/2.2) if self.camera_name=='Canon_EOS_5D' else np.power(16383, 1/2.2) | |
| input_raw_img = np.power(input_raw_img, 1/2.2) | |
| else: | |
| norm_value = 4095 if self.camera_name=='Canon_EOS_5D' else 16383 | |
| target_rgb_img = self.norm_img(target_rgb_img, max_value=255) | |
| input_raw_img = self.norm_img(input_raw_img, max_value=norm_value) | |
| target_raw_img = input_raw_img.copy() | |
| input_raw_img = self.np2tensor(input_raw_img).float() | |
| target_rgb_img = self.np2tensor(target_rgb_img).float() | |
| target_raw_img = self.np2tensor(target_raw_img).float() | |
| sample = {'input_raw':input_raw_img, 'target_rgb':target_rgb_img, 'target_raw':target_raw_img, | |
| 'file_name':input_raw_wb_path.split("/")[-1].split(".")[0]} | |
| return sample | |
| class FiveKDatasetTest(BaseDataset): | |
| def __init__(self, opt): | |
| super().__init__(opt=opt) | |
| self.patch_size = 256 | |
| input_RAWs_WBs, target_RGBs = self.load(is_train=False) | |
| assert len(input_RAWs_WBs) == len(target_RGBs) | |
| self.data = {'input_RAWs_WBs':input_RAWs_WBs, 'target_RGBs':target_RGBs} | |
| def __len__(self): | |
| return len(self.data['input_RAWs_WBs']) | |
| def __getitem__(self, idx): | |
| input_raw_wb_path = self.data['input_RAWs_WBs'][idx] | |
| target_rgb_path = self.data['target_RGBs'][idx] | |
| target_rgb_img = imread(target_rgb_path) | |
| input_raw_wb = np.load(input_raw_wb_path) | |
| input_raw_img = input_raw_wb['raw'] | |
| wb = input_raw_wb['wb'] | |
| wb = wb / wb.max() | |
| input_raw_img = input_raw_img * wb[:-1] | |
| if self.gamma: | |
| norm_value = np.power(4095, 1/2.2) if self.camera_name=='Canon_EOS_5D' else np.power(16383, 1/2.2) | |
| input_raw_img = np.power(input_raw_img, 1/2.2) | |
| else: | |
| norm_value = 4095 if self.camera_name=='Canon_EOS_5D' else 16383 | |
| target_rgb_img = self.norm_img(target_rgb_img, max_value=255) | |
| input_raw_img = self.norm_img(input_raw_img, max_value=norm_value) | |
| target_raw_img = input_raw_img.copy() | |
| input_raw_img = self.np2tensor(input_raw_img).float() | |
| target_rgb_img = self.np2tensor(target_rgb_img).float() | |
| target_raw_img = self.np2tensor(target_raw_img).float() | |
| sample = {'input_raw':input_raw_img, 'target_rgb':target_rgb_img, 'target_raw':target_raw_img, | |
| 'file_name':input_raw_wb_path.split("/")[-1].split(".")[0]} | |
| return sample | |