Merge branch 'add-fargo-dataset' into 'master'

Add fargo dataset

See merge request !10

bob/learn/pytorch/datasets/__init__.py

@@ -4,6 +4,7 @@ from .data_folder import DataFolder
 
 # transforms
 from .utils import FaceCropper
+from .utils import FaceCropAlign
 from .utils import RollChannels 
 from .utils import ToTensor 
 from .utils import Normalize 

bob/learn/pytorch/datasets/fargo.py

+#!/usr/bin/env python
+# encoding: utf-8
+
+import os
+import numpy
+
+from torch.utils.data import Dataset, DataLoader
+
+import bob.db.fargo
+import bob.io.base
+import bob.io.image
+
+from bob.extension import rc
+from bob.db.base import read_annotation_file
+
+from .utils import map_labels
+
+
+class FargoDataset(Dataset):
+  """Class representing the FARGO dataset
+
+  Only retrieves the RGB training set
+
+  Attributes
+  ----------
+  original_directory : str
+    The path to the data
+  transform : `torchvision.transforms`
+    The transform(s) to apply to the face images
+  data_files : list of :obj:`str`
+    The list of data files
+  id_labels : list of :obj:`int`
+    The list of identities, for each data file
+  annotations : 
+    The annotations (eyes center) corresponding to each file 
+
+  """
+  def __init__(self, original_directory=rc['bob.db.fargo.directory'],
+                     annotation_directory=rc['bob.db.fargo.annotation_directory'],
+                     transform=None, start_index=0):
+    """Init function
+
+    Parameters
+    ----------
+    original_directory : str
+      The path to the data
+    annotation_directory : str
+      The path to the annotations
+    transform : :py:class:`torchvision.transforms`
+      The transform(s) to apply to the face images
+    start_index : int
+      label of the first identity (useful if you use several databases)
+    
+    """
+    self.transform = transform
+    self.data_files = []
+    self.annotations = []
+    id_labels = []
+
+    db = bob.db.fargo.Database(original_directory=original_directory)
+    objs = db.objects(purposes='train', modality='rgb')
+ 
+    for o in objs:
+      self.data_files.append(o.make_path(directory=original_directory, extension='.png'))
+      id_labels.append(o.client_id)
+      annotation_file = os.path.join(annotation_directory, o.path + '.pos')
+      self.annotations.append(read_annotation_file(annotation_file, 'eyecenter'))
+
+    self.id_labels = map_labels(id_labels)
+
+  def __len__(self):
+    """Returns the length of the dataset (i.e. nb of examples)
+    
+    Returns
+    -------
+    int 
+      the number of examples in the dataset
+    
+    """
+    return len(self.data_files)
+
+
+  def __getitem__(self, idx):
+    """Returns a sample from the dataset
+   
+    Returns
+    -------
+    dict
+      an example of the dataset, containing the 
+      transformed face image and its identity
+    
+    """
+    image = bob.io.base.load(self.data_files[idx])
+    eyescenter = self.annotations[idx]
+    identity = self.id_labels[idx]
+    sample = {'image': image, 'label': identity, 'eyes' : eyescenter}
+
+    if self.transform:
+      sample = self.transform(sample)
+
+    return sample

bob/learn/pytorch/datasets/utils.py

@@ -9,22 +9,77 @@ class FaceCropper():
   """
     Class to crop a face, based on eyes position
   """
-  def __init__(self, cropped_height, cropped_width):
+  def __init__(self, cropped_height, cropped_width, color_channel='rgb'):
     # the face cropper
     from bob.bio.face.preprocessor import FaceCrop
     cropped_image_size = (cropped_height, cropped_width)
     right_eye_pos = (cropped_height // 5, cropped_width // 4 -1)
     left_eye_pos = (cropped_height // 5, cropped_width // 4 * 3)
     cropped_positions = {'leye': left_eye_pos, 'reye': right_eye_pos}
+    self.color_channel = color_channel
     self.face_cropper = FaceCrop(cropped_image_size=cropped_image_size,
                                  cropped_positions=cropped_positions,
-                                 color_channel='rgb',
+                                 color_channel=color_channel,
                                  dtype='uint8'
                                 )
 
   def __call__(self, sample):
     cropped = self.face_cropper(sample['image'], sample['eyes'])
     sample['image'] = cropped
+    if self.color_channel == 'gray': 
+      sample['image'] = sample['image'][..., numpy.newaxis]
+    return sample
+
+
+class FaceCropAlign():
+  """
+    Wrapper to the FaceCropAlign of bob.pad.face preprocessor
+  
+  """
+
+  def __init__(self, face_size, rgb_output_flag=False,
+               use_face_alignment=True,
+               alignment_type='lightcnn',
+               face_detection_method='mtcnn',
+               ):
+    """ Init function
+
+    Parameters
+    ----------
+
+    face_size: :obj:`int`
+      The size of the cropped face (square)
+    rgb_output_flag: :py:class:`bool`
+      Return RGB cropped face if True, grayscale otherwise 
+    use_face_alignment: :py:class:`bool`
+      If set to True, the face will be aligned, using the facial landmarks detected locally
+      Works only when ``face_detection_method is not None``.
+    alignment_type: :py:class:`str`
+      Specifies the alignment type to use if ``use_face_alignment`` is set to ``True``.
+      Two methods are currently implemented:
+      ``default`` which would do alignment by making eyes horizontally
+      ``lightcnn`` which aligns the face such that eye center and mouth centers are aligned to
+      predefined positions. This option overrides the face size option as the output required
+      is always 128x128. This is suitable for use with LightCNN model.
+    face_detection_method: :py:class:`str`
+      A package to be used for face detection and landmark detection.
+      Options supported by this class: "dlib" and "mtcnn"
+    
+    """
+    from bob.pad.face.preprocessor import FaceCropAlign
+    self.face_cropper = FaceCropAlign(face_size, 
+                                      rgb_output_flag,
+                                      use_face_alignment,
+                                      alignment_type=alignment_type,
+                                      face_detection_method=face_detection_method,
+                                      )
+
+  def __call__(self, sample):
+    cropped = self.face_cropper(sample['image'])
+    if cropped is None:
+      print("Face not detected ...")
+      cropped = numpy.zeros((128, 128))
+    sample['image'] = cropped[..., numpy.newaxis]
     return sample
 
 

bob/learn/pytorch/trainers/CNNTrainer.py

@@ -70,35 +70,36 @@ class CNNTrainer(object):
     """
     try:
       cp = torch.load(model_filename)
-      #self.network.load_state_dict(cp['state_dict'])
     except RuntimeError:
       # pre-trained model was probably saved using nn.DataParallel ...
       cp = torch.load(model_filename, map_location='cpu')
-      if 'state_dict' in cp:
-        from collections import OrderedDict
-        new_state_dict = OrderedDict()
-        for k, v in cp['state_dict'].items():
-          name = k[7:]
-          new_state_dict[name] = v
-        cp['state_dict'] = new_state_dict
-
-    print(type(self.network))
+    
+    if 'state_dict' in cp:
+      from collections import OrderedDict
+      new_state_dict = OrderedDict()
+      for k, v in cp['state_dict'].items():
+        name = k[7:]
+        new_state_dict[name] = v
+    cp['state_dict'] = new_state_dict
 
     ###########################################################################################################
     ### for each defined architecture, get the output size in pre-trained model, and change it if necessary ###
 
     # LightCNN9
     if isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN9):
-      
-      num_classes_pretrained = cp['state_dict']['fc2.weight'].shape[0]
+     
+      last_layer_weight = 'fc2.weight'
+      last_layer_bias = 'fc2.bias'
+
+      num_classes_pretrained = cp['state_dict'][last_layer_weight].shape[0]
       
       if num_classes_pretrained == self.num_classes:
         self.network.load_state_dict(cp['state_dict'])
       else:
-        var = 1.0 / (cp['state_dict']['fc2.weight'].shape[0])
-        np_weights =  numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict']['fc2.weight'].shape[1]))
-        cp['state_dict']['fc2.weight'] = torch.from_numpy(np_weights)
-        cp['state_dict']['fc2.bias'] = torch.zeros(((self.num_classes+1),))
+        var = 1.0 / (cp['state_dict'][last_layer_weight].shape[0])
+        np_weights =  numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict'][last_layer_weight].shape[1]))
+        cp['state_dict'][last_layer_weight] = torch.from_numpy(np_weights)
+        cp['state_dict'][last_layer_bias] = torch.zeros(((self.num_classes+1),))
         #self.network.load_state_dict(cp['state_dict'], strict=False)
         self.network.load_state_dict(cp['state_dict'], strict=True)
 
@@ -197,7 +198,7 @@ class CNNTrainer(object):
       
       start_epoch, start_iter, losses = self.load_and_initialize_model(model)
       if start_epoch != 0: 
-        logger.info('Previous network was trained up to epoch {}, iteration {}'.format(start_epoch, start_iter, losses[-1]))
+        logger.info('Previous network was trained up to epoch {}, iteration {}'.format(start_epoch, start_iter))
         if losses:
           logger.info('Last loss = {}'.format(losses[-1]))
       else: