next

natanielruiz

2017-08-07 6f71fb102f509d705d3abaa1f44638a19f57e92e

 code/datasets.py

@@ -84,20 +84,59 @@

        # 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)
        # And convert to degrees.
        pitch = pose[0] * 180 / np.pi
        yaw = pose[1] * 180 / np.pi
        roll = pose[2] * 180 / np.pi
        # Bin values
        bins = np.array(range(-99, 102, 3))
        labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1)

        if self.transform is not None:
            img = self.transform(img)

        return img, label, self.X_train[index]
        return img, labels, self.X_train[index]

    def __len__(self):
        # 122,450
        return self.length

class AFLW2000_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 degrees.
        pitch = pose[0] * 180 / np.pi
        yaw = pose[1] * 180 / np.pi
        roll = pose[2] * 180 / np.pi
        # Bin values
        bins = np.array(range(-99, 102, 3))
        labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1)

        if self.transform is not None:
            img = self.transform(img)

        return img, labels, self.X_train[index]

    def __len__(self):
        # 2,000
        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