[datasets] added frontal only option (to test with DC-GAN) and changed the RollChannels transform

bob/learn/pytorch/datasets/multipie.py

@@ -16,7 +16,7 @@ import bob.io.image
 class MultiPIEDataset(Dataset):
   """MultiPIE dataset."""
 
-  def __init__(self, root_dir, world=False, transform=None):
+  def __init__(self, root_dir, world=False, frontal_only=False, transform=None):
     """
     """
     self.root_dir = root_dir
@@ -78,6 +78,10 @@ class MultiPIEDataset(Dataset):
       if (camera == '19_1') or (camera == '08_1'):
         continue
 
+      # skip cameras that are not frontal if we want frontal only
+      if (camera != '05_1') and frontal_only:
+        continue
+
       for obj in objs:
         temp = os.path.split(obj.path)
         identity = int(temp[0].split('/')[2])
@@ -88,7 +92,6 @@ class MultiPIEDataset(Dataset):
         self.pose_labels.append(camera_to_label[camera])
       
     self.id_labels = self.map_labels(id_labels)
-    print max(self.id_labels)
       
     
   def __len__(self):
@@ -102,7 +105,7 @@ class MultiPIEDataset(Dataset):
       sample = {'image': image, 'id': identity, 'pose': pose}
 
       if self.transform:
-        sample = self.transform(sample)
+        sample = self.transform(sample['image'])
 
       return sample
 
@@ -121,3 +124,12 @@ class MultiPIEDataset(Dataset):
     
     return labels
 
+class RollChannels(object):
+
+  def __init__(self):
+    pass
+
+  def __call__(self, sample):
+    temp = numpy.rollaxis(numpy.rollaxis(sample, 2),2)
+    sample = temp
+    return sample