-
-
-
class ChaojieDataset(Dataset):
-
def __init__(self,label_list,transforms=None,train=True,test=False):
-
self.test = test
-
self.train = train
-
imgs = []
-
if self.test:
-
for index,row in label_list.iterrows():
-
imgs.append((row["filename"]))
-
self.imgs = imgs
-
else:
-
for index,row in label_list.iterrows():
-
imgs.append((row["filename"],row["label"]))
-
self.imgs = imgs
-
if transforms is None:
-
if self.test or not train:
-
self.transforms = T.Compose([
-
T.Resize((config.img_weight,config.img_height)),
-
T.ToTensor(),
-
T.Normalize(mean = [0.485,0.456,0.406],
-
std = [0.229,0.224,0.225])])
-
else:
-
self.transforms = T.Compose([
-
T.Resize((config.img_weight,config.img_height)),
-
T.RandomRotation(30),
-
T.RandomHorizontalFlip(),
-
T.RandomVerticalFlip(),
-
T.RandomAffine(45),
-
Contrast(1.8),
-
T.ToTensor(),
-
T.Normalize(mean = [0.485,0.456,0.406],
-
std = [0.229,0.224,0.225])])
-
else:
-
self.transforms = transforms
-
def __getitem__(self,index):
-
if self.test:
-
filename = self.imgs[index]
-
img = Image.open(filename)
-
img = self.transforms(img)
-
return img,filename
-
else:
-
filename,label = self.imgs[index]
-
img = Image.open(filename)
-
img = self.transforms(img)
-
return img,label
-
def __len__(self):
-
return len(self.imgs)