Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """ | |
| Created on Wed Feb 12 20:00:46 2020 | |
| @author: Administrator | |
| """ | |
| import os | |
| import os.path | |
| import random | |
| import numpy as np | |
| import cv2 | |
| import h5py | |
| import torch | |
| import torch.utils.data as udata | |
| import argparse | |
| from PIL import Image | |
| class Dataset(udata.Dataset): | |
| r"""Implements torch.utils.data.Dataset | |
| """ | |
| def __init__(self, file, trainrgb=True,trainsyn = True, shuffle=False): | |
| super(Dataset, self).__init__() | |
| self.trainrgb = trainrgb | |
| self.trainsyn = trainsyn | |
| self.train_haze = file | |
| h5f = h5py.File(self.train_haze, 'r') | |
| self.keys = list(h5f.keys()) | |
| if shuffle: | |
| random.shuffle(self.keys) | |
| h5f.close() | |
| def __len__(self): | |
| return len(self.keys) | |
| def __getitem__(self, index): | |
| h5f = h5py.File(self.train_haze, 'r') | |
| key = self.keys[index] | |
| data = np.array(h5f[key]) | |
| h5f.close() | |
| return torch.Tensor(data) | |
| def data_augmentation(clear, mode): | |
| r"""Performs dat augmentation of the input image | |
| Args: | |
| image: a cv2 (OpenCV) image | |
| mode: int. Choice of transformation to apply to the image | |
| 0 - no transformation | |
| 1 - flip up and down | |
| 2 - rotate counterwise 90 degree | |
| 3 - rotate 90 degree and flip up and down | |
| 4 - rotate 180 degree | |
| 5 - rotate 180 degree and flip | |
| 6 - rotate 270 degree | |
| 7 - rotate 270 degree and flip | |
| """ | |
| clear = np.transpose(clear, (2, 3, 0, 1)) | |
| if mode == 0: | |
| # original | |
| clear = clear | |
| elif mode == 1: | |
| # flip up and down | |
| clear = np.flipud(clear) | |
| elif mode == 2: | |
| # rotate counterwise 90 degree | |
| clear = np.rot90(clear) | |
| elif mode == 3: | |
| # rotate 90 degree and flip up and down | |
| clear = np.rot90(clear) | |
| clear = np.flipud(clear) | |
| elif mode == 4: | |
| # rotate 180 degree | |
| clear = np.rot90(clear, k=2) | |
| elif mode == 5: | |
| # rotate 180 degree and flip | |
| clear = np.rot90(clear, k=2) | |
| clear = np.flipud(clear) | |
| elif mode == 6: | |
| # rotate 270 degree | |
| clear = np.rot90(clear, k=3) | |
| elif mode == 7: | |
| # rotate 270 degree and flip | |
| clear = np.rot90(clear, k=3) | |
| clear = np.flipud(clear) | |
| else: | |
| raise Exception('Invalid choice of image transformation') | |
| return np.transpose(clear, (2, 3, 0, 1)) | |
| def img_to_patches(img,win,stride,Syn=True): | |
| typ, chl, raw, col = img.shape | |
| chl = int(chl) | |
| num_raw = np.ceil((raw-win)/stride+1).astype(np.uint8) | |
| num_col = np.ceil((col-win)/stride+1).astype(np.uint8) | |
| count = 0 | |
| total_process = int(num_col)*int(num_raw) | |
| img_patches = np.zeros([typ, chl, win, win, total_process]) | |
| if Syn: | |
| for i in range(num_raw): | |
| for j in range(num_col): | |
| if stride * i + win <= raw and stride * j + win <=col: | |
| img_patches[:,:,:,:,count] = img[:, :, stride*i : stride*i + win, stride*j : stride*j + win] | |
| elif stride * i + win > raw and stride * j + win<=col: | |
| img_patches[:,:,:,:,count] = img[:, :,raw-win : raw,stride * j : stride * j + win] | |
| elif stride * i + win <= raw and stride*j + win>col: | |
| img_patches[:,:,:,:,count] = img[:, :,stride*i : stride*i + win, col-win : col] | |
| else: | |
| img_patches[:,:,:,:,count] = img[:, :,raw-win : raw,col-win : col] | |
| img_patches[:,:,:,:,count] = data_augmentation(img_patches[:, :, :, :, count], np.random.randint(0, 7)) | |
| count +=1 | |
| return img_patches | |
| def read_img(img): | |
| return np.array(Image.open(img))/255. | |
| def Train_data(args): | |
| file_list = os.listdir(f'{args.train_path}/{args.gt_name}') | |
| with h5py.File(args.data_name, 'w') as h5f: | |
| count = 0 | |
| for i in range(len(file_list)): | |
| print(file_list[i]) | |
| img_list = [] | |
| img_list.append(read_img(f'{args.train_path}/{args.gt_name}/{file_list[i]}')) | |
| for j in args.degradation_name: | |
| img_list.append(read_img(f'{args.train_path}/{j}/{file_list[i]}')) | |
| img = np.stack(img_list,0) | |
| img = img_to_patches(img.transpose(0, 3, 1, 2), args.patch_size, args.stride) | |
| for nx in range(img.shape[4]): | |
| data = img[:,:,:,:,nx] | |
| print(count, data.shape) | |
| h5f.create_dataset(str(count), data=data) | |
| count += 1 | |
| h5f.close() | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description = "Building the training patch database") | |
| parser.add_argument("--patch-size", type = int, default=256, help="Patch size") | |
| parser.add_argument("--stride", type = int, default=200, help="Size of stride") | |
| parser.add_argument("--train-path", type = str, default='./data/CDD-11_train', help="Train path") | |
| parser.add_argument("--data-name", type = str, default='dataset.h5', help="Data name") | |
| parser.add_argument("--gt-name", type = str, default='clear', help="HQ name") | |
| parser.add_argument("--degradation-name", type = list, default=['low','haze','rain','snow',\ | |
| 'low_haze','low_rain','low_snow','haze_rain','haze_snow','low_haze_rain','low_haze_snow'], help="LQ name") | |
| args = parser.parse_args() | |
| Train_data(args) |