-
Notifications
You must be signed in to change notification settings - Fork 17
/
data.py
114 lines (97 loc) · 4.3 KB
/
data.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
111
112
from torch.utils.data import Dataset
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
from math import floor
from utils import landmarks_68_to_5
def process(img , landmarks_5pts):
batch = {}
name = ['left_eye','right_eye','nose','mouth']
patch_size = {
'left_eye':(40,40),
'right_eye':(40,40),
'nose':(40,32),
'mouth':(48,32),
}
landmarks_5pts[3,0] = (landmarks_5pts[3,0] + landmarks_5pts[4,0]) / 2.0
landmarks_5pts[3,1] = (landmarks_5pts[3,1] + landmarks_5pts[4,1]) / 2.0
# crops
for i in range(4):
x = floor(landmarks_5pts[i,0])
y = floor(landmarks_5pts[i,1])
batch[ name[i] ] = img.crop( (x - patch_size[ name[i] ][0]//2 + 1 , y - patch_size[ name[i] ][1]//2 + 1 , x + patch_size[ name[i] ][0]//2 + 1 , y + patch_size[ name[i] ][1]//2 + 1 ) )
return batch
class TrainDataset( Dataset):
def __init__( self , img_list ):
super(type(self),self).__init__()
self.img_list = img_list
def __len__( self ):
return len(self.img_list)
def __getitem__( self , idx ):
#filename processing
batch = {}
img_name = self.img_list[idx].split('/')
img_frontal_name = self.img_list[idx].split('_')
img_frontal_name[-2] = '051'
img_frontal_name = '_'.join( img_frontal_name ).split('/')
batch['img'] = Image.open( '/'.join( img_name ) )
batch['img32'] = Image.open( '/'.join( img_name[:-2] + ['32x32' , img_name[-1] ] ) )
batch['img64'] = Image.open( '/'.join( img_name[:-2] + ['64x64' , img_name[-1] ] ) )
batch['img_frontal'] = Image.open( '/'.join(img_frontal_name) )
batch['img32_frontal'] = Image.open( '/'.join( img_frontal_name[:-2] + ['32x32' , img_frontal_name[-1] ] ) )
batch['img64_frontal'] = Image.open( '/'.join( img_frontal_name[:-2] + ['64x64' , img_frontal_name[-1] ] ) )
patch_name_list = ['left_eye','right_eye','nose','mouth']
for patch_name in patch_name_list:
batch[patch_name] = Image.open( '/'.join(img_name[:-2] + ['patch' , patch_name , img_name[-1] ]) )
batch[patch_name+'_frontal'] = Image.open( '/'.join(img_frontal_name[:-2] + ['patch' , patch_name , img_frontal_name[-1] ]) )
totensor = transforms.ToTensor()
for k in batch:
batch[k] = totensor( batch[k] )
batch[k] = batch[k] *2.0 -1.0
#if batch[k].max() <= 0.9:
# print( "{} {} {}".format( batch[k].max(), self.img_list[idx] , k ))
#if batch[k].min() >= -0.9:
# print( "{} {} {}".format( batch[k].min() , self.img_list[idx] , k ) )
batch['label'] = int( self.img_list[idx].split('/')[-1].split('_')[0] )
return batch
class PretrainDataset( Dataset):
def __init__( self , img_list ):
super(type(self),self).__init__()
self.img_list = img_list
def __len__( self):
return len(self.img_list)
def __getitem__(self,idx):
batch = {}
totensor = transforms.ToTensor()
img = Image.open( self.img_list[idx] )
img = totensor( img )
img = img*2.0 - 1.0
batch['img'] = img
batch['label'] = int( self.img_list[idx].split('/')[-1].split('_')[0] )
return batch
class TestDataset( Dataset):
def __init__( self , img_list , lm_list):
super(type(self),self).__init__()
self.img_list = img_list
self.lm_list = lm_list
assert len(img_list) == len(lm_list)
def __len__(self):
return len(self.img_list)
def __getitem__(self,idx):
img_name = self.img_list[idx]
img = Image.open( img_name )
lm = np.array( self.lm_list[idx].split(' ') , np.float32 ).reshape(-1,2)
lm = landmarks_68_to_5( lm )
for i in range(5):
lm[i][0] *= 128/img.width
lm[i][1] *= 128/img.height
img = img.resize( (128,128) , Image.LANCZOS)
batch = process( img , lm )
batch['img'] = img
batch['img64'] = img.resize( (64,64) , Image.LANCZOS )
batch['img32'] = batch['img64'].resize( (32,32) , Image.LANCZOS )
to_tensor = transforms.ToTensor()
for k in batch:
batch[k] = to_tensor( batch[k] )
batch[k] = batch[k] * 2.0 - 1.0
return batch