1 import torch 2 import os,glob 3 import random,csv 4 5 from torch.utils.data import Dataset,DataLoader 6 7 from torchvision import transforms 8 from PIL import Image 9 10 class Pokemon(Dataset): 11 ''' 12 @param 13 root:存储的根路径 14 resize：将图片大小根据网络结构适配 15 mode：train或者test模式 16 ''' 17 def __init__(self,root,resize,mode): 18 super(Pokemon,self).__init__() 19 20 self.root = root 21 self.resize = resize 22 23 # 字典类型key:name value:label 24 self.name2label = {} 25 # listdir返回顺序不固定，用sorted将它固定，因为排序一次之后就固定了 26 for name in sorted(os.listdir(os.path.join(root))): 27 if not os.path.isdir(os.path.join(root,name)): 28 continue 29 30 self.name2label[name] = len(self.name2label.keys()) 31 32 # print(self.name2label) 33 34 # image_path + image_label 35 self.images,self.labels = self.load_csv('images.csv') 36 37 if mode == 'train': # 60% 38 self.images = self.images[:int(0.6*len(self.images))] 39 self.labels = self.labels[:int(0.6*len(self.labels))] 40 elif mode == 'val': # 20% 41 self.images = self.images[int(0.6*len(self.images)):int(0.8*len(self.images))] 42 self.labels = self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))] 43 elif mode == 'test': # 20% = 80% ->100% 44 self.images = self.images[int(0.8*len(self.images)):] 45 self.labels = self.labels[int(0.8*len(self.labels)):] 46 47 def load_csv(self,filename): 48 49 # 如果不存在再写入，存在的话直接读取就可以了 50 if not os.path.exists(os.path.join(self.root,filename)) 51 images = [] 52 for name in self.name2label.keys(): 53 # 'pokemon'\mewtwo\00001.png 54 images += glob.glob(os.path.join(self.root,name,'*.png')) 55 images += glob.glob(os.path.join(self.root,name,'*.jpg')) 56 images += glob.glob(os.path.join(self.root,name,'*.jpeg')) 57 58 # 1167，'pokemon\bulbasaur\00000000.png' 59 print(len(images),images) 60 61 random.shuffle(images) 62 with open(os.path.join(self.root,filename),mode = 'w',newline='') as f: 63 writer = csv.writer(f) 64 for img in images: # 'pokemon\bulbasaur\00000000.png' 65 name = img.split(os.sep)[-2] 66 label = self.name2label[name] 67 # 'pokemon\bulbasaur\00000000.png',0 68 writer.writerow([img,label]) 69 print('writen into csv file:',filename) 70 71 # read from csv file 72 images,labels = [],[] 73 with open(os.path.join(self.root,filename)) 74 reader = csv.reader(f) 75 for row in reader: 76 # 'pokemon\bulbasaur\00000000.png',0 77 img,label = row 78 label = int(label) 79 80 images.append(img) 81 labels.append(label) 82 83 # 保证images和labels一一对应，长度相等 84 assert len(images) == len(labels) 85 return images,labels 86 87 def __len__(self): 88 89 return len(self.images) 90 91 def denormalize(self,x_hat)： 92 93 mean=[0.485,0.456,0.406] 94 std=[0.229,0.224,0.225] 95 96 # x_hat = (x-mean)/std 97 # x = x_hat*std+mean 98 # x: [c,h,w] 99 # mean: [3] --> [3,1,1] 100 mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1) 101 std = torch.tensor(std).unsqueeze(1).unsqueeze(1) 102 103 x = x_hat*std + mean 104 105 return x 106 107 108 def __getitem__(self,idx): 109 # idx~[0~len(images)] 110 # self.images,self.labels 111 # img: pokemon\bulbasaur\00000000.png' 112 # label: 0 113 img,label = self.images[idx],self.labels[idx] 114 115 tf = transforms.Compose([ 116 lambda x:Image.open(x).convert('RGB'), # string path --> image data 117 transforms.Resize((int(self.resize*1.25),int(self.resize*1.25))), 118 transforms.RandomRotation(15), 119 transforms.CenterCrop(self.resize), 120 transforms.ToTensor(), 121 transforms.Normalize(mean=[0.485,0.456,0.406], 122 std=[0.229,0.224,0.225]) 123 ]) 124 125 img = tf(img) 126 label = torch.tensor(label) 127 128 return img,label