From 31fc66b795c0a57b8009d7b03f49f6cd099ceb29 Mon Sep 17 00:00:00 2001 From: natanielruiz <nataniel777@hotmail.com> Date: 星期六, 23 九月 2017 12:07:48 +0800 Subject: [PATCH] Trying superres --- code/datasets.py | 65 +++++++++++++++++++++++++++++++- 1 files changed, 62 insertions(+), 3 deletions(-) diff --git a/code/datasets.py b/code/datasets.py index da0603f..17f1899 100644 --- a/code/datasets.py +++ b/code/datasets.py @@ -106,15 +106,16 @@ # 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))) @@ -138,6 +139,64 @@ # 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) + # 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 AFLW_aug(Dataset): def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.txt', image_mode='RGB'): self.data_dir = data_dir -- Gitblit v1.8.0