From b74b9c54247177c82493f180617d7551de8e2bb1 Mon Sep 17 00:00:00 2001 From: natanielruiz <nataniel777@hotmail.com> Date: 星期二, 26 九月 2017 03:47:06 +0800 Subject: [PATCH] Before SR experiments --- code/datasets.py | 74 +++++++++++++++++++++++++++++++----- 1 files changed, 63 insertions(+), 11 deletions(-) diff --git a/code/datasets.py b/code/datasets.py index 5825105..17f1899 100644 --- a/code/datasets.py +++ b/code/datasets.py @@ -106,17 +106,76 @@ # 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.15 - x_min -= 0.6 * k * abs(x_max - x_min) + k = 0.20 + x_min -= 2 * 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) + 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))) + + # 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_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))) + + ds = 5 + 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) # We get the pose in radians pose = utils.get_ypr_from_mat(mat_path) @@ -353,7 +412,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 @@ -364,13 +423,6 @@ 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))) - - # Flip? - # rnd = np.random.random_sample() - # if rnd < 0.5: - # yaw = -yaw - # roll = -roll - # img = img.transpose(Image.FLIP_LEFT_RIGHT) # Bin values bins = np.array(range(-99, 102, 3)) -- Gitblit v1.8.0