[datasets] made some modif to be able to concatenate datasets

bob/learn/pytorch/datasets/__init__.py

@@ -7,6 +7,7 @@ from .utils import ToTensor
 from .utils import Normalize 
 
 from .utils import map_labels
+from .utils import ConcatDataset
 
 
 # gets sphinx autodoc done right - don't remove it

bob/learn/pytorch/datasets/casia_webface.py

@@ -33,7 +33,7 @@ class CasiaDataset(Dataset):
 
   # TODO: Start from original data and annotations - Guillaume HEUSCH, 06-11-2017
 
-  def __init__(self, root_dir, frontal_only=False, transform=None):
+  def __init__(self, root_dir, frontal_only=False, transform=None, start_index=0):
     self.root_dir = root_dir
     self.transform = transform
   
@@ -65,9 +65,9 @@ class CasiaDataset(Dataset):
         cluster = path[-2]
         self.data_files.append(os.path.join(root, name))
         self.pose_labels.append(int(dir_to_pose_label[cluster]))
-        id_labels.append(int(subject))
+        id_labels.append(subject)
  
-    self.id_labels = map_labels(id_labels)
+    self.id_labels = map_labels(id_labels, start_index)
       
     
   def __len__(self):

bob/learn/pytorch/datasets/multipie.py

@@ -103,7 +103,7 @@ class MultiPIEDataset(Dataset):
       for obj in objs:
         temp = os.path.split(obj.path)
         identity = int(temp[0].split('/')[2])
-        id_labels.append(int(identity))
+        id_labels.append(identity)
         cropped_filename = os.path.join(root_dir, camera_to_pose[camera], temp[1])
         cropped_filename += '.png'
         self.data_files.append(cropped_filename)

bob/learn/pytorch/datasets/utils.py

@@ -32,18 +32,63 @@ class Normalize(object):
     sample['image'] = self.op(sample['image'])
     return sample
 
-def map_labels(raw_labels):
+def map_labels(raw_labels, start_index=0):
   """
   Map the ID label to [0 - # of IDs] 
   """
-  
   possible_labels = list(set(raw_labels))
-
   labels = numpy.array(raw_labels)
   
   for i in range(len(possible_labels)):
     l = possible_labels[i]
-    labels[numpy.where(labels==l)[0]] = i
-  
+    labels[numpy.where(labels==l)[0]] = i + start_index
+
   return labels
 
+
+from torch.utils.data import Dataset
+import bob.io.base
+import bob.io.image
+
+class ConcatDataset(Dataset):
+  """
+  Class to concatenate two or more datasets for DR-GAN training
+
+  **Parameters**
+
+  datasets: list
+    The list of datasets (as torch.utils.data.Dataset)
+  """
+  def __init__(self, datasets):
+    
+    self.transform = datasets[0].transform
+    self.data_files = sum((d.data_files for d in datasets), [])
+    self.pose_labels = sum((d.pose_labels for d in datasets), [])
+  
+    # TODO: for an unknown reason, the above does not work for id_labels - Guillaume HEUSCH, 01-12-2017
+    self.id_labels = []
+    for d in datasets:
+      self.id_labels.append(d.id_labels)
+      print "type of id labels = {}".format(type(d.id_labels[0]))
+    print len(self.id_labels)
+
+  def __len__(self):
+      """
+        return the length of the dataset (i.e. nb of examples)
+      """
+      return len(self.data_files)
+
+
+  def __getitem__(self, idx):
+      """
+        return a sample from the dataset
+      """
+      image = bob.io.base.load(self.data_files[idx])
+      identity = self.id_labels[idx]
+      pose = self.pose_labels[idx]
+      sample = {'image': image, 'id': identity, 'pose': pose}
+
+      if self.transform:
+        sample = self.transform(sample)
+
+      return sample