GoodWin/Deep-Multi-scale

1# -- coding: utf-8 --2import os.path3import random4import torchvision.transforms as transforms5import torch6from data.base_dataset import BaseDataset7from data.image_folder import make_dataset8from PIL import Image, ImageFilter9import numpy as np10import cv211import math12from util import util13from scipy.io import loadmat14from PIL import Image15import PIL16 17 18class AlignedDataset(BaseDataset):19    @staticmethod20    def modify_commandline_options(parser, is_train):21        return parser22 23    def initialize(self, opt):24        self.opt = opt25        self.root = opt.dataroot26        self.partpath = opt.partroot27        self.dir_AB = os.path.join(opt.dataroot, opt.phase)28        self.AB_paths = sorted(make_dataset(self.dir_AB))29        self.is_real = opt.is_real30        # assert(opt.resize_or_crop == 'resize_and_crop')31        assert(opt.resize_or_crop == 'degradation')   32 33    def AddNoise(self,img): # noise34        if random.random() > 0.9: #35            return img36        self.sigma = np.random.randint(1, 11)37        img_tensor = torch.from_numpy(np.array(img)).float()38        noise = torch.randn(img_tensor.size()).mul_(self.sigma/1.0)39 40        noiseimg = torch.clamp(noise+img_tensor,0,255)41        return Image.fromarray(np.uint8(noiseimg.numpy()))42 43    def AddBlur(self,img): # gaussian blur or motion blur44        if random.random() > 0.9: #45            return img46        img = np.array(img)47        if random.random() > 0.35: ##gaussian blur48            blursize = random.randint(1,17) * 2 + 1 ##3,5,7,9,11,13,1549            blursigma = random.randint(3, 20)50            img = cv2.GaussianBlur(img, (blursize,blursize), blursigma/10)51        else: #motion blur52            M = random.randint(1,32)53            KName = './data/MotionBlurKernel/m_%02d.mat' % M54            k = loadmat(KName)['kernel']55            k = k.astype(np.float32)56            k /= np.sum(k)57            img = cv2.filter2D(img,-1,k)58        return Image.fromarray(img)59 60    def AddDownSample(self,img): # downsampling61        if random.random() > 0.95: #62            return img63        sampler = random.randint(20, 100)*1.064        img = img.resize((int(self.opt.fineSize/sampler*10.0), int(self.opt.fineSize/sampler*10.0)), Image.BICUBIC)65        return img66 67    def AddJPEG(self,img): # JPEG compression68        if random.random() > 0.6: #69            return img70        imQ = random.randint(40, 80)71        img = np.array(img)72        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY),imQ] # (0,100),higher is better,default is 9573        _, encA = cv2.imencode('.jpg',img,encode_param)74        img = cv2.imdecode(encA,1)75        return Image.fromarray(img)76 77    def AddUpSample(self,img):78        return img.resize((self.opt.fineSize, self.opt.fineSize), Image.BICUBIC)79 80    def __getitem__(self, index): # 81 82        AB_path = self.AB_paths[index]83        Imgs = Image.open(AB_path).convert('RGB')84        # # 85        A = Imgs.resize((self.opt.fineSize, self.opt.fineSize))86        A = transforms.ColorJitter(0.3, 0.3, 0.3, 0)(A)87        C = A88        A = self.AddUpSample(self.AddJPEG(self.AddNoise(self.AddDownSample(self.AddBlur(A)))))89 90        tmps = AB_path.split('/')91        ImgName = tmps[-1]92        Part_locations = self.get_part_location(self.partpath, ImgName, 2)93   94        A = transforms.ToTensor()(A) # 95        C = transforms.ToTensor()(C)96        97        ##98        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(A) # 99        C = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(C) # 100        return {'A':A, 'C':C, 'A_paths': AB_path,'Part_locations': Part_locations}101 102    def get_part_location(self, landmarkpath, imgname, downscale=1):103        Landmarks = []104        with open(os.path.join(landmarkpath,imgname+'.txt'),'r') as f:105            for line in f:106                tmp = [np.float(i) for i in line.split(' ') if i != '\n']107                Landmarks.append(tmp)108        Landmarks = np.array(Landmarks)/downscale # 512 * 512109        110        Map_LE = list(np.hstack((range(17,22), range(36,42))))111        Map_RE = list(np.hstack((range(22,27), range(42,48))))112        Map_NO = list(range(29,36))113        Map_MO = list(range(48,68))114        #left eye115        Mean_LE = np.mean(Landmarks[Map_LE],0)116        L_LE = np.max((np.max(np.max(Landmarks[Map_LE],0) - np.min(Landmarks[Map_LE],0))/2,16))117        Location_LE = np.hstack((Mean_LE - L_LE + 1, Mean_LE + L_LE)).astype(int)118        #right eye119        Mean_RE = np.mean(Landmarks[Map_RE],0)120        L_RE = np.max((np.max(np.max(Landmarks[Map_RE],0) - np.min(Landmarks[Map_RE],0))/2,16))121        Location_RE = np.hstack((Mean_RE - L_RE + 1, Mean_RE + L_RE)).astype(int)122        #nose123        Mean_NO = np.mean(Landmarks[Map_NO],0)124        L_NO = np.max((np.max(np.max(Landmarks[Map_NO],0) - np.min(Landmarks[Map_NO],0))/2,16))125        Location_NO = np.hstack((Mean_NO - L_NO + 1, Mean_NO + L_NO)).astype(int)126        #mouth127        Mean_MO = np.mean(Landmarks[Map_MO],0)128        L_MO = np.max((np.max(np.max(Landmarks[Map_MO],0) - np.min(Landmarks[Map_MO],0))/2,16))129 130        Location_MO = np.hstack((Mean_MO - L_MO + 1, Mean_MO + L_MO)).astype(int)131        return Location_LE, Location_RE, Location_NO, Location_MO132 133    def __len__(self): #134        return len(self.AB_paths)135 136    def name(self):137        return 'AlignedDataset'138