From 6664c6d52fad58e396861946a3bed7d5afc4d44d Mon Sep 17 00:00:00 2001
From: natanielruiz <nataniel777@hotmail.com>
Date: 星期五, 07 七月 2017 10:53:52 +0800
Subject: [PATCH] Training for hopenet works.
---
code/datasets.py | 49 +++++++++++++++++++++++++++++++++++++++++++++----
1 files changed, 45 insertions(+), 4 deletions(-)
diff --git a/code/datasets.py b/code/datasets.py
index 030059f..3750e71 100644
--- a/code/datasets.py
+++ b/code/datasets.py
@@ -84,10 +84,51 @@
# 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
+ # 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)
- label = torch.FloatTensor(pose)
+ if self.transform is not None:
+ img = self.transform(img)
+
+ 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, yaw, roll = pose * 180 / np.pi
+ # Bin values
+ bins = np.array(range(-99, 102, 3))
+ binned_pitch = torch.DoubleTensor(np.digitize(pitch, bins) - 1)
+ binned_yaw = torch.DoubleTensor(np.digitize(yaw, bins) - 1)
+ binned_roll = torch.DoubleTensor(np.digitize(roll, bins) - 1)
+
+ label = binned_yaw, binned_pitch, binned_roll
if self.transform is not None:
img = self.transform(img)
@@ -95,7 +136,7 @@
return img, label, self.X_train[index]
def __len__(self):
- # 122,450
+ # 2,000
return self.length
def get_list_from_filenames(file_path):
--
Gitblit v1.8.0