Start adding packnet model

packnet_functional.py

@@ -0,0 +1,98 @@
+import tensorflow as tf
+import tensorflow.keras as keras
+import tensorflow.keras.layers as layers
+from tensorflow import nn
+
+import group_norm
+
+
+def pack_layer():
+    pass
+
+
+def residual_layer(inputs, out_channels, stride, dropout=None):
+    """
+    Keras implementation of the Residual block (ResNet) as used in PackNet
+    :param inputs:
+    :param out_channels:
+    :param stride:
+    :param dropout:
+    :return:
+    """
+    x = layers.Conv2D(out_channels, 3, padding='same', stride=stride)(inputs)
+    x = layers.Conv2D(out_channels, 3, padding='same')(x)
+    shortcut = layers.Conv2D(out_channels, 3, padding='same', stride=stride)(inputs)
+    if dropout:
+        shortcut = keras.layers.SpatialDropout2D(dropout)(shortcut)
+    x = keras.layers.Concatenate()([x, shortcut])
+    x = group_norm.GroupNormalization(16)(x)
+    return keras.layers.ELU()(x)
+
+
+# Packnet usually expects more than one layer per block (2,2,3,3)
+def residual_bock(inputs, out_channels, residual_layers, stride, dropout=None):
+    pass
+
+
+def packnet_conv2d(inputs, out_channels, kernel_size, stride):
+    x = keras.layers.Conv2D(out_channels, kernel_size, stride, padding='same')
+    x = group_norm.GroupNormalization(16)(x)
+    return keras.layers.ELU()(x)
+
+
+def packnet_inverse_depth(inputs, out_channels=1, min_depth=0.5):
+    x = packnet_conv2d(inputs, out_channels, kernel_size=3, stride=1)
+    return keras.activations.sigmoid(x) / min_depth
+
+
+def pack_3d(inputs, kernel_size, r=2, features_3d=8):
+    """
+    Implementatino of the 3d packing block proposed here: https://arxiv.org/abs/1905.02693
+    :param inputs:
+    :param kernel_size:
+    :param r:
+    :param features_3d:
+    :return:
+    """
+    # Data format for single image in nyu is HWC (space_to_depth uses NHWC as default)
+    x = nn.space_to_depth(inputs, r)
+    x = tf.expand_dims(x, 1)
+    x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')
+    b, c, d, h, w = x.shape
+    x = tf.reshape(x, (b, c * d, h, w))
+    return packnet_conv2d(x, inputs.shape[1], kernel_size, 1)
+
+
+def unpack_3d(inputs, out_channels, kernel_size, r=3, features_3d=8):
+    x = packnet_conv2d(inputs, out_channels * (r ** 2) // features_3d, kernel_size, 1)
+    x = tf.expand_dims(x, 1)  # B x D x 4(out)/D x H/2 x W/2
+    x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')
+    b, c, d, h, w = x.shape
+    x = tf.reshape(x, [b, c * d, h, w])
+    return nn.depth_to_space(x, r)
+
+
+# TODO: Support different size packnet for scaling up/down
+def make_packnet(shape=(224, 224, 3), skip_add=True, features_3d=4):
+    """
+    Make the PackNet depth network.
+    :param shape: Input shape of the image
+    :param skip_add: Set to use add rather than concat skip connections, defaults to True
+    :return:
+    """
+    input = keras.layers.Input(shape=shape)
+    x = packnet_conv2d(input, 32, 5, 1)
+    skip_1 = x
+    x = packnet_conv2d(input, 32, 7, 1)
+    x = pack_3d(x, 5, features_3d)
+    x = residual_layer(x, 64, )
+
+    # TODO: Skip connection
+    if skip_add:
+        x = keras.layers.Add([x, ])
+    else:
+        x = keras.layers.Concatenate([x, ])
+
+    x = packnet_conv2d(x, 32, 3, 1)
+    x = packnet_inverse_depth(x)
+    return keras.Model(inputs=input, outputs=x, name="PackNet")