Spaces:
Build error
Build error
| from random import randint | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from torch.utils.data.dataset import Dataset | |
| from utils import cvtColor, preprocess_input | |
| from torch.utils.data import DataLoader | |
| def get_new_img_size(width, height, img_min_side=600): | |
| if width <= height: | |
| f = float(img_min_side) / width | |
| resized_height = int(f * height) | |
| resized_width = int(img_min_side) | |
| else: | |
| f = float(img_min_side) / height | |
| resized_width = int(f * width) | |
| resized_height = int(img_min_side) | |
| return resized_width, resized_height | |
| class SRGANDataset(Dataset): | |
| def __init__(self, train_lines, lr_shape, hr_shape): | |
| super(SRGANDataset, self).__init__() | |
| self.train_lines = train_lines | |
| self.train_batches = len(train_lines) | |
| self.lr_shape = lr_shape | |
| self.hr_shape = hr_shape | |
| def __len__(self): | |
| return self.train_batches | |
| def __getitem__(self, index): | |
| index = index % self.train_batches | |
| image_origin = Image.open(self.train_lines[index].split()[0]) | |
| if self.rand()<.5: | |
| img_h = self.get_random_data(image_origin, self.hr_shape) | |
| else: | |
| img_h = self.random_crop(image_origin, self.hr_shape[1], self.hr_shape[0]) | |
| img_l = img_h.resize((self.lr_shape[1], self.lr_shape[0]), Image.BICUBIC) | |
| img_h = np.transpose(preprocess_input(np.array(img_h, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]) | |
| img_l = np.transpose(preprocess_input(np.array(img_l, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]) | |
| return np.array(img_l), np.array(img_h) | |
| def rand(self, a=0, b=1): | |
| return np.random.rand()*(b-a) + a | |
| def get_random_data(self, image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True): | |
| #------------------------------# | |
| # 读取图像并转换成RGB图像 | |
| #------------------------------# | |
| image = cvtColor(image) | |
| #------------------------------# | |
| # 获得图像的高宽与目标高宽 | |
| #------------------------------# | |
| iw, ih = image.size | |
| h, w = input_shape | |
| if not random: | |
| scale = min(w/iw, h/ih) | |
| nw = int(iw*scale) | |
| nh = int(ih*scale) | |
| dx = (w-nw)//2 | |
| dy = (h-nh)//2 | |
| #---------------------------------# | |
| # 将图像多余的部分加上灰条 | |
| #---------------------------------# | |
| image = image.resize((nw,nh), Image.BICUBIC) | |
| new_image = Image.new('RGB', (w,h), (128,128,128)) | |
| new_image.paste(image, (dx, dy)) | |
| image_data = np.array(new_image, np.float32) | |
| return image_data | |
| #------------------------------------------# | |
| # 对图像进行缩放并且进行长和宽的扭曲 | |
| #------------------------------------------# | |
| new_ar = w/h * self.rand(1-jitter,1+jitter)/self.rand(1-jitter,1+jitter) | |
| scale = self.rand(1, 1.5) | |
| if new_ar < 1: | |
| nh = int(scale*h) | |
| nw = int(nh*new_ar) | |
| else: | |
| nw = int(scale*w) | |
| nh = int(nw/new_ar) | |
| image = image.resize((nw,nh), Image.BICUBIC) | |
| #------------------------------------------# | |
| # 将图像多余的部分加上灰条 | |
| #------------------------------------------# | |
| dx = int(self.rand(0, w-nw)) | |
| dy = int(self.rand(0, h-nh)) | |
| new_image = Image.new('RGB', (w,h), (128,128,128)) | |
| new_image.paste(image, (dx, dy)) | |
| image = new_image | |
| #------------------------------------------# | |
| # 翻转图像 | |
| #------------------------------------------# | |
| flip = self.rand()<.5 | |
| if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT) | |
| rotate = self.rand()<.5 | |
| if rotate: | |
| angle = np.random.randint(-15,15) | |
| a,b = w/2,h/2 | |
| M = cv2.getRotationMatrix2D((a,b),angle,1) | |
| image = cv2.warpAffine(np.array(image), M, (w,h), borderValue=[128,128,128]) | |
| #------------------------------------------# | |
| # 色域扭曲 | |
| #------------------------------------------# | |
| hue = self.rand(-hue, hue) | |
| sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat) | |
| val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val) | |
| x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV) | |
| x[..., 1] *= sat | |
| x[..., 2] *= val | |
| x[x[:,:, 0]>360, 0] = 360 | |
| x[:, :, 1:][x[:, :, 1:]>1] = 1 | |
| x[x<0] = 0 | |
| image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255 | |
| return Image.fromarray(np.uint8(image_data)) | |
| def random_crop(self, image, width, height): | |
| #--------------------------------------------# | |
| # 如果图像过小无法截取，先对图像进行放大 | |
| #--------------------------------------------# | |
| if image.size[0] < self.hr_shape[1] or image.size[1] < self.hr_shape[0]: | |
| resized_width, resized_height = get_new_img_size(width, height, img_min_side=np.max(self.hr_shape)) | |
| image = image.resize((resized_width, resized_height), Image.BICUBIC) | |
| #--------------------------------------------# | |
| # 随机截取一部分 | |
| #--------------------------------------------# | |
| width1 = randint(0, image.size[0] - width) | |
| height1 = randint(0, image.size[1] - height) | |
| width2 = width1 + width | |
| height2 = height1 + height | |
| image = image.crop((width1, height1, width2, height2)) | |
| return image | |
| def SRGAN_dataset_collate(batch): | |
| images_l = [] | |
| images_h = [] | |
| for img_l, img_h in batch: | |
| images_l.append(img_l) | |
| images_h.append(img_h) | |
| return np.array(images_l), np.array(images_h) | |