add posed mesh, metrics, tanh, l1 with gt

stylegan2/layers/bias_act.py

@@ -7,12 +7,15 @@ import tensorflow as tf
 class BiasAct(tf.keras.layers.Layer):
     def __init__(self, lrmul, act, **kwargs):
         super(BiasAct, self).__init__(**kwargs)
-        assert act in ['linear', 'lrelu']
+        assert act in ['linear', 'lrelu', 'tanh']
         self.lrmul = lrmul
 
         if act == 'linear':
             self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x))
             self.gain = 1.0
+        elif act == 'tanh':
+            self.act = tf.keras.layers.Activation('tanh')
+            self.gain = 1.0
         else:
             self.act = tf.keras.layers.LeakyReLU(alpha=0.2)
             self.gain = np.sqrt(2)

stylegan2/layers/synthesis_block.py

@@ -3,7 +3,7 @@ import tensorflow as tf
 from stylegan2.layers.modulated_conv2d import ModulatedConv2D
 from stylegan2.layers.bias_act import BiasAct
 from stylegan2.layers.noise import Noise
-from stylegan2.layers.to_rgb import ToRGB
+from stylegan2.layers.to_rgb import ToRGB, ToRGBLast
 from stylegan2.layers.cuda.upfirdn_2d_v2 import upsample_2d, compute_paddings
 
 
@@ -120,12 +120,17 @@ class Synthesis(tf.keras.layers.Layer):
         prev_n_f = n_f
         self.blocks = list()
         self.torgbs = list()
-        for res, n_f in zip(self.resolutions[1:], self.featuremaps[1:]):
+        for res, n_f in zip(self.resolutions[1:-1], self.featuremaps[1:-1]):
             self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res,
                                               name='{:d}x{:d}/block'.format(res, res)))
             self.torgbs.append(ToRGB(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res)))
             prev_n_f = n_f
 
+        res, n_f = self.resolutions[-1], self.featuremaps[-1]
+        self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res,
+                                            name='{:d}x{:d}/block'.format(res, res)))
+        self.torgbs.append(ToRGBLast(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res)))
+
     def call(self, inputs, training=None, mask=None):
         w_broadcasted = inputs
 

stylegan2/layers/to_rgb.py

@@ -28,3 +28,29 @@ class ToRGB(tf.keras.layers.Layer):
             'res': self.res,
         })
         return config
+
+
+class ToRGBLast(tf.keras.layers.Layer):
+    def __init__(self, in_ch, res, **kwargs):
+        super(ToRGBLast, self).__init__(**kwargs)
+        self.in_ch = in_ch
+        self.res = res
+
+        self.conv = ModulatedConv2D(in_res=res, in_fmaps=in_ch, fmaps=3, kernel=1, up=False, down=False, demodulate=False,
+                                    resample_kernel=None, gain=1.0, lrmul=1.0, fused_modconv=True, name='conv')
+        self.apply_bias = BiasAct(lrmul=1.0, act='tanh', name='bias')
+
+    def call(self, inputs, training=None, mask=None):
+        x, w = inputs
+
+        x = self.conv([x, w])
+        x = self.apply_bias(x)
+        return x
+
+    def get_config(self):
+        config = super(ToRGBLast, self).get_config()
+        config.update({
+            'in_ch': self.in_ch,
+            'res': self.res,
+        })
+        return config