132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
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_block(inputs, out_channels, residual_layers, stride, dropout=None):
|
|
x = inputs
|
|
for i in range(0, residual_layers):
|
|
x = residual_layer(x, out_channels, stride, dropout)
|
|
return x
|
|
|
|
|
|
def packnet_conv2d(inputs, out_channels, kernel_size, stride):
|
|
x = keras.layers.Conv2D(out_channels, kernel_size,
|
|
stride, padding='same')(inputs)
|
|
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, 4)
|
|
x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')(x)
|
|
b, h, w, c, d = x.shape
|
|
x = tf.reshape(x, (b, h, w, c * d))
|
|
return packnet_conv2d(x, inputs.shape[3], kernel_size, 1)
|
|
|
|
|
|
def unpack_3d(inputs, out_channels, kernel_size, r=2, features_3d=8):
|
|
x = packnet_conv2d(inputs, out_channels * (r ** 2) //
|
|
features_3d, kernel_size, 1)
|
|
x = tf.expand_dims(x, 4) # B x H/2 x W/2 x 4(out)/D x D
|
|
x = keras.layers.Conv3D(features_3d, kernel_size=3, padding='same')(x)
|
|
b, h, w, c, d = x.shape
|
|
x = tf.reshape(x, [b, h, w, c * d])
|
|
return nn.depth_to_space(x, r)
|
|
|
|
|
|
# TODO: Support different size packnet for scaling up/down
|
|
# TODO: Support different channel format (right now we're supporting NHWC, we should also support NCHW)
|
|
def make_packnet(shape=(224, 224, 3), skip_add=True, features_3d=4, dropout=None):
|
|
"""
|
|
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:
|
|
"""
|
|
|
|
# ================ ENCODER =================
|
|
input = keras.layers.Input(shape=shape)
|
|
x = packnet_conv2d(input, 32, 5, 1)
|
|
skip_1 = x
|
|
x = packnet_conv2d(x, 32, 7, 1)
|
|
x = pack_3d(x, 5, features_3d=features_3d)
|
|
skip_2 = x
|
|
x = residual_block(x, 64, 2, 1, dropout)
|
|
x = pack_3d(x, 3, features_3d=features_3d)
|
|
skip_3 = x
|
|
x = residual_block(x, 128, 2, 1, dropout)
|
|
x = pack_3d(x, 3, features_3d=features_3d)
|
|
skip_4 = x
|
|
x = residual_block(x, 256, 3, 1, dropout)
|
|
x = pack_3d(x, 3, features_3d=features_3d)
|
|
skip_5 = x
|
|
x = residual_block(x, 512, 3, 1, dropout)
|
|
x = pack_3d(x, 3, features_3d=features_3d)
|
|
# ================ ENCODER =================
|
|
|
|
x = unpack_3d(x, 512, 3, features_3d=features_3d)
|
|
x = keras.layers.Add(
|
|
[x, skip_5]) if skip_add else keras.layers.Concatenate([x, skip_5])
|
|
x = packnet_conv2d(x, 512, 3, 1)
|
|
x = unpack_3d(x, 256, 3, features_3d=features_3d)
|
|
x = keras.layers.Add(
|
|
[x, skip_4]) if skip_add else keras.layers.Concatenate([x, skip_4])
|
|
x = packnet_conv2d(x, 256, 3, 1)
|
|
# TODO: This is wrong, look at the paper
|
|
x = packnet_inverse_depth(x, 1)
|
|
x = keras.layers.UpSampling2D()
|
|
|
|
# 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")
|