-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_loader.py
111 lines (91 loc) · 4.26 KB
/
data_loader.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import math
import numpy as np
import torchvision.models as models
import torch.utils.data as data
from torchvision import transforms
import cv2
import torch.nn.functional as F
from torch.autograd import Variable
import pandas as pd
import os ,torch
import torch.nn as nn
import image_utils
import argparse,random
from torchvision import datasets
class RafDataSet(data.Dataset):
def __init__(self, raf_path, phase, transform=None, basic_aug=False):
self.phase = phase
self.transform = transform
self.raf_path = raf_path
NAME_COLUMN = 0
LABEL_COLUMN = 1
df = pd.read_csv(os.path.join(self.raf_path, 'EmoLabel/new_20_noise.txt'), sep=' ', header=None)
#df = pd.read_csv(os.path.join(self.raf_path, 'EmoLabel/new_10_noise.txt'), sep=' ', header=None)
#df = pd.read_csv(os.path.join(self.raf_path, 'EmoLabel/list_patition_label.txt'), sep=' ', header=None)
if phase == 'train':
dataset = df[df[NAME_COLUMN].str.startswith('train')]
else:
dataset = df[df[NAME_COLUMN].str.startswith('test')]
file_names = dataset.iloc[:, NAME_COLUMN].values
self.label = dataset.iloc[:,
LABEL_COLUMN].values - 1 # 0:Surprise, 1:Fear, 2:Disgust, 3:Happiness, 4:Sadness, 5:Anger, 6:Neutral
self.file_paths = []
# use raf aligned images for training/testing
for f in file_names:
f = f.split(".")[0]
f = f + "_aligned.jpg"
path = os.path.join(self.raf_path, 'Image/aligned', f)
self.file_paths.append(path)
self.basic_aug = basic_aug
self.aug_func = [image_utils.flip_image, image_utils.add_gaussian_noise]
def __len__(self):
return len(self.file_paths)
def __getitem__(self, idx):
path = self.file_paths[idx]
image = cv2.imread(path)
image = image[:, :, ::-1] # BGR to RGB
label = self.label[idx]
# augmentation
if self.phase == 'train':
if self.basic_aug and random.uniform(0, 1) > 0.5:
index = random.randint(0, 1)
image = self.aug_func[index](image)
if self.transform is not None:
image = self.transform(image)
return image, label, idx
def get_train_loader(raf_path, batch_size, num_workers=1, shuffle=True, pin_memory=True):
data_transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((48, 48)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
transforms.RandomErasing(scale=(0.02, 0.25))])
train_dataset = RafDataSet(raf_path, phase='train', transform=data_transforms, basic_aug=True)
print('Train set size:', train_dataset.__len__())
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle,
pin_memory=pin_memory)
return train_loader
def get_valid_loader(raf_path, batch_size, num_workers=1, shuffle=False, pin_memory=True):
data_transforms_val = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((48, 48)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
val_dataset = RafDataSet(raf_path, phase='test', transform=data_transforms_val)
print('Validation set size:', val_dataset.__len__())
valid_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle,
pin_memory=pin_memory)
return valid_loader
if __name__ == '__main__':
get_train_loader('/home/dengshuchao/datasets/RafDb/raf-basic/',64,1,True,True)
print("The test of training dataset is ok!")
get_valid_loader('/home/dengshuchao/datasets/RafDb/raf-basic/', 64, 1, False, True)
print("The test of test dataset is ok!")