Merge branch 'drgan-msu' into 'master'

added the DR-GAN encoder as a bob.bio.base.extractor

See merge request !3

bob/ip/tensorflow_extractor/DrGanMSU.py

+#!/usr/bin/env python
+# encoding: utf-8
+
+
+import numpy
+import tensorflow as tf
+import os
+import bob.ip.base
+import pkg_resources
+from . import download_file
+import logging
+logger = logging.getLogger(__name__)
+
+
+class batch_norm(object):
+    """Code modification of http://stackoverflow.com/a/33950177"""
+
+    def __init__(self, epsilon=1e-5, momentum=0.9, name="batch_norm"):
+        with tf.variable_scope(name):
+            self.epsilon = epsilon
+            self.momentum = momentum
+            self.name = name
+
+    def __call__(self, x, train=True, reuse=False):
+        return tf.contrib.layers.batch_norm(x,
+                                            decay=self.momentum,
+                                            updates_collections=None,
+                                            epsilon=self.epsilon,
+                                            scale=True,
+                                            reuse=reuse,
+                                            is_training=train,
+                                            scope=self.name)
+
+
+def conv2d(input_, output_dim,
+           k_h=3, k_w=3, d_h=2, d_w=2, stddev=0.02,
+           name="conv2d", reuse=False):
+    with tf.variable_scope(name, reuse=reuse):
+        w = tf.get_variable('w', [k_h, k_w, input_.get_shape()[-1], output_dim],
+                            initializer=tf.truncated_normal_initializer(stddev=stddev))
+        conv = tf.nn.conv2d(
+            input_, w, strides=[
+                1, d_h, d_w, 1], padding='SAME')
+
+        biases = tf.get_variable(
+            'biases',
+            [output_dim],
+            initializer=tf.constant_initializer(0.0))
+        conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())
+
+        return conv
+
+
+def elu(x, name='elu'):
+    return tf.nn.elu(x, name)
+
+
+class DRGAN(object):
+    """
+    class implementing the DR GAN, as described in
+
+
+    Note: this code has orginally been done by L.Tran @ MSU, and is heavily influenced
+          by the DCGAN code here:
+
+          I simplified it to the maximum to increase readability, and to limit its usage
+          to face image encoding.
+
+    **Parameters**
+
+    image_size: int
+      The size of the (squared) image.
+
+    gf_dim: int
+      The number of feature maps in the first convolutional layer (encoder and discriminator).
+
+    gfc_dim: int
+      The dimension of the encoded id (output of the encoder).
+    """
+
+    def __init__(self, image_size=96, gf_dim=32, gfc_dim=320):
+
+        self.image_size = image_size
+        self.gf_dim = gf_dim
+        self.gfc_dim = gfc_dim
+
+        # batch normalization
+        self.g_bn0_0 = batch_norm(name='g_k_bn0_0')
+        self.g_bn0_1 = batch_norm(name='g_k_bn0_1')
+        self.g_bn0_2 = batch_norm(name='g_k_bn0_2')
+        self.g_bn0_3 = batch_norm(name='g_k_bn0_3')
+        self.g_bn1_0 = batch_norm(name='g_k_bn1_0')
+        self.g_bn1_1 = batch_norm(name='g_k_bn1_1')
+        self.g_bn1_2 = batch_norm(name='g_k_bn1_2')
+        self.g_bn1_3 = batch_norm(name='g_k_bn1_3')
+        self.g_bn2_0 = batch_norm(name='g_k_bn2_0')
+        self.g_bn2_1 = batch_norm(name='g_k_bn2_1')
+        self.g_bn2_2 = batch_norm(name='g_k_bn2_2')
+        self.g_bn2_3 = batch_norm(name='g_k_bn2_3')
+        self.g_bn3_0 = batch_norm(name='g_k_bn3_0')
+        self.g_bn3_1 = batch_norm(name='g_k_bn3_1')
+        self.g_bn3_2 = batch_norm(name='g_k_bn3_2')
+        self.g_bn3_3 = batch_norm(name='g_k_bn3_3')
+        self.g_bn4_0 = batch_norm(name='g_k_bn4_0')
+        self.g_bn4_1 = batch_norm(name='g_k_bn4_1')
+        self.g_bn4_2 = batch_norm(name='g_k_bn4_2')
+        self.g_bn4_c = batch_norm(name='g_h_bn4_c')
+        self.g_bn5 = batch_norm(name='g_k_bn5')
+
+    def generator_encoder(self, image, is_reuse=False, is_training=True):
+        """ Function that define the graph doing the encoding of a face image.
+
+        **Parameters**
+
+        image: numpy array
+          The input image
+
+        is_reuse: bool
+          Reuse variables.
+
+        is_training: bool
+          Flag for training mode
+
+        **Returns**
+          The encoded id
+        """
+        s16 = int(self.image_size / 16)
+        k0_0 = image
+
+        k0_1 = elu(
+            self.g_bn0_1(
+                conv2d(
+                    k0_0,
+                    self.gf_dim * 1,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k01_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k01_prelu')
+        k0_2 = elu(
+            self.g_bn0_2(
+                conv2d(
+                    k0_1,
+                    self.gf_dim * 2,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k02_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k02_prelu')
+
+        k1_0 = elu(
+            self.g_bn1_0(
+                conv2d(
+                    k0_2,
+                    self.gf_dim * 2,
+                    d_h=2,
+                    d_w=2,
+                    name='g_k10_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k10_prelu')
+        k1_1 = elu(
+            self.g_bn1_1(
+                conv2d(
+                    k1_0,
+                    self.gf_dim * 2,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k11_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k11_prelu')
+        k1_2 = elu(
+            self.g_bn1_2(
+                conv2d(
+                    k1_1,
+                    self.gf_dim * 4,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k12_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k12_prelu')
+
+        k2_0 = elu(
+            self.g_bn2_0(
+                conv2d(
+                    k1_2,
+                    self.gf_dim * 4,
+                    d_h=2,
+                    d_w=2,
+                    name='g_k20_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k20_prelu')
+        k2_1 = elu(
+            self.g_bn2_1(
+                conv2d(
+                    k2_0,
+                    self.gf_dim * 3,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k21_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k21_prelu')
+        k2_2 = elu(
+            self.g_bn2_2(
+                conv2d(
+                    k2_1,
+                    self.gf_dim * 6,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k22_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k22_prelu')
+
+        k3_0 = elu(
+            self.g_bn3_0(
+                conv2d(
+                    k2_2,
+                    self.gf_dim * 6,
+                    d_h=2,
+                    d_w=2,
+                    name='g_k30_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k30_prelu')
+        k3_1 = elu(
+            self.g_bn3_1(
+                conv2d(
+                    k3_0,
+                    self.gf_dim * 4,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k31_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k31_prelu')
+        k3_2 = elu(
+            self.g_bn3_2(
+                conv2d(
+                    k3_1,
+                    self.gf_dim * 8,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k32_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k32_prelu')
+
+        k4_0 = elu(
+            self.g_bn4_0(
+                conv2d(
+                    k3_2,
+                    self.gf_dim * 8,
+                    d_h=2,
+                    d_w=2,
+                    name='g_k40_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k40_prelu')
+        k4_1 = elu(
+            self.g_bn4_1(
+                conv2d(
+                    k4_0,
+                    self.gf_dim * 5,
+                    d_h=1,
+                    d_w=1,
+                    name='g_k41_conv',
+                    reuse=is_reuse),
+                train=is_training,
+                reuse=is_reuse),
+            name='g_k41_prelu')
+        k4_2 = self.g_bn4_2(
+            conv2d(
+                k4_1,
+                self.gfc_dim,
+                d_h=1,
+                d_w=1,
+                name='g_k42_conv',
+                reuse=is_reuse),
+            train=is_training,
+            reuse=is_reuse)
+
+        k5 = tf.nn.avg_pool(
+            k4_2, ksize=[
+                1, s16, s16, 1], strides=[
+                1, 1, 1, 1], padding='VALID')
+        k5 = tf.reshape(k5, [-1, self.gfc_dim])
+
+        # dropout if training
+        if (is_training):
+            k5 = tf.nn.dropout(k5, keep_prob=0.6)
+
+        return k5
+
+
+class DrGanMSUExtractor(object):
+    """Wrapper for the free DRGan by L.Tran @ MSU:
+
+    To use this class as a bob.bio.base extractor::
+
+        from bob.bio.base.extractor import Extractor
+        class DrGanMSUExtractorBioBase(DrGanMSUExtractor, Extractor):
+            pass
+        extractor = DrGanMSUExtractorBioBase()
+
+
+    **Parameters:**
+
+      model_file:
+        Path to the model
+        
+      image_size: list
+        The input image size (WxHxC)
+      
+    """
+
+    def __init__(self, model_path=None, image_size=[96, 96, 3]):
+
+        self.image_size = image_size
+        self.session = tf.Session()
+
+        # placeholder for the input image
+        data_shape = [1] + self.image_size
+        self.X = tf.placeholder(tf.float32, shape=data_shape)
+
+        # the encoder
+        self.drgan = DRGAN()
+        self.encode = self.drgan.generator_encoder(
+            self.X, is_reuse=False, is_training=False)
+
+        # If the path is not, set the default path
+        if model_path is None:
+            model_path = self.get_modelpath()
+
+        # If does not exist, download
+        if not os.path.exists(model_path):
+            self.download_model()
+                
+        self.saver = tf.train.Saver()
+        # Reestore either from the last checkpoint or from a particular checkpoint
+        if os.path.isdir(model_path):
+            self.saver.restore(self.session,
+                               tf.train.latest_checkpoint(model_path))
+        else:
+            self.saver.restore(self.session, model_path)
+
+    @staticmethod
+    def get_modelpath():
+        import pkg_resources
+        return pkg_resources.resource_filename(__name__,
+                                               'data/DR_GAN_model')
+
+    @staticmethod
+    def download_model():
+        """
+        Download and extract the DrGanMSU files in bob/ip/tensorflow_extractor
+        """
+        import zipfile
+        zip_file = os.path.join(DrGanMSUExtractor.get_modelpath(),
+                                "DR_GAN_model.zip")
+        urls = [
+            # This is a private link at Idiap to save bandwidth.
+            "http://beatubulatest.lab.idiap.ch/private/wheels/gitlab/"
+            "DR_GAN_model.zip",
+        ]
+
+        for url in urls:
+            try:
+                logger.info(
+                    "Downloading the DrGanMSU model from "
+                    "{} ...".format(url))
+                download_file(url, zip_file)
+                break
+            except Exception:
+                logger.warning(
+                    "Could not download from the %s url", url, exc_info=True)
+        else:  # else is for the for loop
+            if not os.path.isfile(zip_file):
+                raise RuntimeError("Could not download the zip file.")
+
+        # Unzip
+        logger.info("Unziping in {0}".format(DrGanMSUExtractor.get_modelpath()))
+        with zipfile.ZipFile(zip_file) as myzip:
+            myzip.extractall(os.path.dirname(DrGanMSUExtractor.get_modelpath()))
+
+        # delete extra files
+        os.unlink(zip_file)
+
+    def __call__(self, image):
+        """__call__(image) -> feature
+
+        Extract features
+
+        **Parameters:**
+
+        image : 3D :py:class:`numpy.ndarray` (floats)
+          The image to extract the features from.
+
+        **Returns:**
+
+        feature : 2D :py:class:`numpy.ndarray` (floats)
+          The extracted features
+        """
+        def bob2skimage(bob_image):
+            """
+            Convert bob color image to the skcit image
+            """
+
+            if len(bob_image.shape) == 3:
+                skimage = numpy.zeros(
+                    shape=(
+                        bob_image.shape[1],
+                        bob_image.shape[2],
+                        3),
+                    dtype=numpy.uint8)
+                skimage[:, :, 0] = bob_image[0, :, :]
+                skimage[:, :, 1] = bob_image[1, :, :]
+                skimage[:, :, 2] = bob_image[2, :, :]
+            else:
+                skimage = numpy.zeros(
+                    shape=(
+                        bob_image.shape[0],
+                        bob_image.shape[1],
+                        1))
+                skimage[:, :, 0] = bob_image[:, :]
+            return skimage
+
+        def rescaleToUint8(image):
+            result = numpy.zeros_like(image)
+            for channel in range(image.shape[2]):
+                min_image = numpy.min(image[:, :, channel])
+                max_image = numpy.max(image[:, :, channel])
+                if (max_image - min_image) != 0:
+                    result[:, :, channel] = 255.0 * \
+                        ((image[:, :, channel] - min_image) / (max_image - min_image))
+                else:
+                    result[:, :, channel] = 0
+                result = result.astype('uint8')
+            return result
+
+        # encode the provided image
+        image = rescaleToUint8(image)
+        image = bob2skimage(image)
+        image = numpy.array(image / 127.5 - 1).astype(numpy.float32)
+        shape = [1] + list(image.shape)
+        img = numpy.reshape(image, tuple(shape))
+        encoded_id = self.session.run(self.encode, feed_dict={self.X: img})
+        return encoded_id
+

bob/ip/tensorflow_extractor/__init__.py

@@ -66,6 +66,7 @@ def get_config():
 
 from .Extractor import Extractor
 from .FaceNet import FaceNet
+from .DrGanMSU import  DrGanMSUExtractor
 
 
 # gets sphinx autodoc done right - don't remove it
@@ -87,6 +88,7 @@ def __appropriate__(*args):
 __appropriate__(
     Extractor,
     FaceNet,
+    DrGanMSUExtractor,
 )
 
 # gets sphinx autodoc done right - don't remove it

bob/ip/tensorflow_extractor/test.py

@@ -46,6 +46,18 @@ def test_facenet():
     output = extractor(data)
     assert output.size == 128, output.shape
 
+def test_drgan():
+    """
+    '/remote/idiap.svm/user.active/heusch/work/dev/DR-GAN_code_wmodel/DR_GAN_model/DCGAN.model-590000'
+    """
+    from bob.ip.tensorflow_extractor import DrGanMSUExtractor
+    #extractor = DrGanMSUExtractor("/idiap/project/hface/models/cnn/DR_GAN_model/", image_size=[96, 96, 3])
+    extractor = DrGanMSUExtractor()
+    data = numpy.random.rand(3, 96, 96).astype("uint8")
+    output = extractor(data)
+    assert output.size == 320, output.shape
+
+
 
 """
 def test_output_from_meta():

doc/extra-intersphinx.txt

+bob.bio.base

doc/guide.rst

@@ -48,3 +48,24 @@ In this example we plot some outputs of the convolutional layer `conv1`.
 
 .. plot:: plot/convolve.py
    :include-source: False
+   
+   
+Facenet Model
+-------------
+
+
+:ref:`bob.bio.base <bob.bio.base>` wrapper Facenet model.
+Check `here for more info <py_api.html#bob.ip.tensorflow_extractor.FaceNet>`_
+
+
+
+DRGan from L.Tran @ MSU:
+------------------------
+
+:ref:`bob.bio.base <bob.bio.base>` wrapper to the DRGan model trained by L.Tran @ MSU.
+Check `here <py_api.html#bob.ip.tensorflow_extractor.DrGanMSUExtractor>`_ for more info
+
+
+
+
+