natanielruiz
2017-09-27 43416c4717d2430c3e11f042294d12b781fee2e1
code/datasets.py
@@ -6,6 +6,7 @@
from PIL import Image, ImageFilter
import utils
from torchvision import transforms
def stack_grayscale_tensor(tensor):
    tensor = torch.cat([tensor, tensor, tensor], 0)
@@ -80,6 +81,168 @@
            img = self.transform(img)
        return img, labels, cont_labels, self.X_train[index]
    def __len__(self):
        # 122,450
        return self.length
class Pose_300W_LP_random_ds(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'):
        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.image_mode = image_mode
        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(self.image_mode)
        mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext)
        shape_path = os.path.join(self.data_dir, self.y_train[index] + '_shape.npy')
        # Crop the face
        pt2d = utils.get_pt2d_from_mat(mat_path)
        x_min = min(pt2d[0,:])
        y_min = min(pt2d[1,:])
        x_max = max(pt2d[0,:])
        y_max = max(pt2d[1,:])
        # k = 0.2 to 0.40
        k = np.random.random_sample() * 0.2 + 0.2
        x_min -= 0.6 * k * abs(x_max - x_min)
        y_min -= 2 * k * abs(y_max - y_min)
        x_max += 0.6 * k * abs(x_max - x_min)
        y_max += 0.6 * k * abs(y_max - y_min)
        img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
        # We get the pose in radians
        pose = utils.get_ypr_from_mat(mat_path)
        # And convert to degrees.
        pitch = pose[0] * 180 / np.pi
        yaw = pose[1] * 180 / np.pi
        roll = pose[2] * 180 / np.pi
        rnd = np.random.random_sample()
        if rnd < 0.5:
            ds = 10
            original_size = img.size
            img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST)
            img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST)
        # Flip?
        rnd = np.random.random_sample()
        if rnd < 0.5:
            yaw = -yaw
            roll = -roll
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
        # Blur?
        rnd = np.random.random_sample()
        if rnd < 0.05:
            img = img.filter(ImageFilter.BLUR)
        # Bin values
        bins = np.array(range(-99, 102, 3))
        binned_pose = np.digitize([yaw, pitch, roll], bins) - 1
        # Get shape
        shape = np.load(shape_path)
        labels = torch.LongTensor(np.concatenate((binned_pose, shape), axis = 0))
        cont_labels = torch.FloatTensor([yaw, pitch, roll])
        if self.transform is not None:
            img = self.transform(img)
        return img, labels, cont_labels, self.X_train[index]
    def __len__(self):
        # 122,450
        return self.length
class Pose_300W_LP_SR(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'):
        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.image_mode = image_mode
        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(self.image_mode)
        mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext)
        # Crop the face
        pt2d = utils.get_pt2d_from_mat(mat_path)
        x_min = min(pt2d[0,:])
        y_min = min(pt2d[1,:])
        x_max = max(pt2d[0,:])
        y_max = max(pt2d[1,:])
        # k = 0.2 to 0.40
        k = np.random.random_sample() * 0.2 + 0.2
        x_min -= 0.6 * k * abs(x_max - x_min)
        y_min -= 2 * k * abs(y_max - y_min)
        x_max += 0.6 * k * abs(x_max - x_min)
        y_max += 0.6 * k * abs(y_max - y_min)
        img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
        # We get the pose in radians
        pose = utils.get_ypr_from_mat(mat_path)
        # And convert to degrees.
        pitch = pose[0] * 180 / np.pi
        yaw = pose[1] * 180 / np.pi
        roll = pose[2] * 180 / np.pi
        rnd = np.random.random_sample()
        if rnd < 0.5:
            ds = 10
            original_size = img.size
            img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST)
            img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST)
        # Flip?
        rnd = np.random.random_sample()
        if rnd < 0.5:
            yaw = -yaw
            roll = -roll
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
        # Blur?
        rnd = np.random.random_sample()
        if rnd < 0.05:
            img = img.filter(ImageFilter.BLUR)
        img_ycc = img.convert('YCbCr')
        # Bin values
        bins = np.array(range(-99, 102, 3))
        binned_pose = np.digitize([yaw, pitch, roll], bins) - 1
        labels = torch.LongTensor(np.concatenate((binned_pose, shape), axis = 0))
        cont_labels = torch.FloatTensor([yaw, pitch, roll])
        # Transforms
        img = transforms.Scale(240)(img)
        img = transforms.RandomCrop(224)(img)
        img_ycc = img.convert('YCbCr')
        img = transforms.ToTensor()
        img_ycc = transforms.ToTensor()
        return img, img_ycc, labels, cont_labels, self.X_train[index]
    def __len__(self):
        # 122,450
@@ -172,10 +335,70 @@
        y_max += 0.6 * k * abs(y_max - y_min)
        img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
        ds = 5
        ds = 8
        original_size = img.size
        img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=0)
        img = img.resize((original_size[0], original_size[1]), resample=0)
        img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST)
        img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST)
        # We get the pose in radians
        pose = utils.get_ypr_from_mat(mat_path)
        # 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)
        cont_labels = torch.FloatTensor([yaw, pitch, roll])
        if self.transform is not None:
            img = self.transform(img)
        return img, labels, cont_labels, self.X_train[index]
    def __len__(self):
        # 2,000
        return self.length
class AFLW2000_random_ds(Dataset):
    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat', image_mode='RGB'):
        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.image_mode = image_mode
        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(self.image_mode)
        mat_path = os.path.join(self.data_dir, self.y_train[index] + self.annot_ext)
        # Crop the face
        pt2d = utils.get_pt2d_from_mat(mat_path)
        x_min = min(pt2d[0,:])
        y_min = min(pt2d[1,:])
        x_max = max(pt2d[0,:])
        y_max = max(pt2d[1,:])
        k = 0.20
        x_min -= 2 * k * abs(x_max - x_min)
        y_min -= 2 * k * abs(y_max - y_min)
        x_max += 2 * k * abs(x_max - x_min)
        y_max += 0.6 * k * abs(y_max - y_min)
        img = img.crop((int(x_min), int(y_min), int(x_max), int(y_max)))
        rnd = np.random.random_sample()
        if rnd < 0.5:
            ds = 10
            original_size = img.size
            img = img.resize((img.size[0] / ds, img.size[1] / ds), resample=Image.NEAREST)
            img = img.resize((original_size[0], original_size[1]), resample=Image.NEAREST)
        # We get the pose in radians
        pose = utils.get_ypr_from_mat(mat_path)
@@ -412,7 +635,7 @@
        R = R[:3,:]
        pose_annot.close()
        roll = -np.arctan2(R[1][0], R[0][0]) * 180 / np.pi
        roll = np.arctan2(R[1][0], R[0][0]) * 180 / np.pi
        yaw = np.arctan2(-R[2][0], np.sqrt(R[2][1] ** 2 + R[2][2] ** 2)) * 180 / np.pi
        pitch = -np.arctan2(R[2][1], R[2][2]) * 180 / np.pi