First commit

natanielruiz

2017-07-07 ce90d92fad8d1e9ad87553ffd8dd9c553cfcc67c

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
import numpy as np
import torch
import cv2
from torch.utils.data.dataset import Dataset
import os
from PIL import Image
 
import utils
 
class Pose_300W_LP(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat'):
        self.data_dir = data_dir
        self.transform = transform
        self.img_ext = img_ext
        self.annot_ext = annot_ext
 
        filename_list = get_list_from_filenames(filename_path)
 
        self.X_train = filename_list
        self.y_train = filename_list
        self.length = len(filename_list)
 
    def __getitem__(self, index):
        img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
        img = img.convert('RGB')
 
        pose = utils.get_ypr_from_mat(os.path.join(self.data_dir, self.y_train[index] + self.annot_ext))
        label = torch.FloatTensor(pose)
 
        if self.transform is not None:
            img = self.transform(img)
 
        return img, label, self.X_train[index]
 
    def __len__(self):
        # 122,450
        return self.length
 
class AFLW2000(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat'):
        self.data_dir = data_dir
        self.transform = transform
        self.img_ext = img_ext
        self.annot_ext = annot_ext
 
        filename_list = get_list_from_filenames(filename_path)
 
        self.X_train = filename_list
        self.y_train = filename_list
        self.length = len(filename_list)
 
    def __getitem__(self, index):
        img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
        img = img.convert('RGB')
 
        pose = utils.get_ypr_from_mat(os.path.join(self.data_dir, self.y_train[index] + self.annot_ext))
        label = torch.FloatTensor(pose)
 
        if self.transform is not None:
            img = self.transform(img)
 
        return img, label, self.X_train[index]
 
    def __len__(self):
        # 2,000
        return self.length
 
class Pose_300W_LP_binned(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat'):
        self.data_dir = data_dir
        self.transform = transform
        self.img_ext = img_ext
        self.annot_ext = annot_ext
 
        filename_list = get_list_from_filenames(filename_path)
 
        self.X_train = filename_list
        self.y_train = filename_list
        self.length = len(filename_list)
 
    def __getitem__(self, index):
        img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
        img = img.convert('RGB')
 
        # We get the pose in radians
        pose = utils.get_ypr_from_mat(os.path.join(self.data_dir, self.y_train[index] + self.annot_ext))
        # And convert to positive degrees.
        pose = pose * 180 / np.pi + 90
 
        label = torch.FloatTensor(pose)
 
        if self.transform is not None:
            img = self.transform(img)
 
        return img, label, self.X_train[index]
 
    def __len__(self):
        # 122,450
        return self.length
 
def get_list_from_filenames(file_path):
    # input:    relative path to .txt file with file names
    # output:   list of relative path names
    with open(file_path) as f:
        lines = f.read().splitlines()
    return lines