# Path: nn_utils/activations.py
# def from_hdr_activation(x):
#     return tf.math.log(1 + tf.nn.relu(x))
# 
# def to_hdr_activation(x):
#     return tf.exp(tf.nn.relu(x)) - 1
# 
# Path: nn_utils/coord_conv.py
# class AddCoords(tf.keras.layers.Layer):
#     """Add coords to a tensor"""
# 
#     def __init__(
#         self,
#     ):
#         super(AddCoords, self).__init__()
# 
#     def call(self, input_tensor):
#         """
#         input_tensor: (batch, 1, 1, c), or (batch, x_dim, y_dim, c)
#         In the first case, first tile the input_tensor to be (batch, x_dim, y_dim, c)
#         In the second case, skiptile, just concat
#         """
#         batch_size_tensor = tf.shape(input_tensor)[0]
#         x_dim = tf.shape(input_tensor)[1]
#         y_dim = tf.shape(input_tensor)[2]
#         xx_ones = tf.ones([batch_size_tensor, x_dim], dtype=tf.int32)
#         xx_ones = tf.expand_dims(xx_ones, -1)
#         xx_range = tf.tile(tf.expand_dims(tf.range(y_dim), 0), [batch_size_tensor, 1])
#         xx_range = tf.expand_dims(xx_range, 1)
#         xx_channel = tf.matmul(xx_ones, xx_range)
#         xx_channel = tf.expand_dims(xx_channel, -1)
#         yy_ones = tf.ones([batch_size_tensor, y_dim], dtype=tf.int32)
#         yy_ones = tf.expand_dims(yy_ones, 1)
#         yy_range = tf.tile(tf.expand_dims(tf.range(x_dim), 0), [batch_size_tensor, 1])
#         yy_range = tf.expand_dims(yy_range, -1)
#         yy_channel = tf.matmul(yy_range, yy_ones)
#         yy_channel = tf.expand_dims(yy_channel, -1)
# 
#         x_dim = tf.cast(x_dim, tf.float32)
#         y_dim = tf.cast(y_dim, tf.float32)
# 
#         xx_channel = tf.cast(xx_channel, tf.float32) / (x_dim - 1)
#         yy_channel = tf.cast(yy_channel, tf.float32) / (y_dim - 1)
#         xx_channel = xx_channel * 2 - 1
#         yy_channel = yy_channel * 2 - 1
# 
#         ret = tf.concat([input_tensor, xx_channel, yy_channel], axis=-1)
# 
#         output_tensor = ret
# 
#         return output_tensor


import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
from nn_utils.activations import from_hdr_activation, to_hdr_activation
from nn_utils.coord_conv import AddCoords
                output_units,
                1,
                padding="same",
                activation=to_hdr_activation,
            )
        )

        self.net = tf.keras.Sequential(net)

    @tf.function
    def call(self, embedding):
        reconstruction = self.net(embedding)

        return reconstruction


class CnnDiscriminator(tf.keras.Model):
    def __init__(
        self,
        discriminator_units: int = 32,
        activation="relu",
        img_height: int = 128,
        **kwargs,
    ):
        super(CnnDiscriminator, self).__init__(**kwargs)

        activation = tf.keras.layers.Activation(activation)

        net = [
            tf.keras.layers.InputLayer((img_height, img_height * 2, 3)),

# Path: nn_utils/activations.py
# def from_hdr_activation(x):
#     return tf.math.log(1 + tf.nn.relu(x))
# 
# def to_hdr_activation(x):
#     return tf.exp(tf.nn.relu(x)) - 1
# 
# Path: nn_utils/coord_conv.py
# class AddCoords(tf.keras.layers.Layer):
#     """Add coords to a tensor"""
# 
#     def __init__(
#         self,
#     ):
#         super(AddCoords, self).__init__()
# 
#     def call(self, input_tensor):
#         """
#         input_tensor: (batch, 1, 1, c), or (batch, x_dim, y_dim, c)
#         In the first case, first tile the input_tensor to be (batch, x_dim, y_dim, c)
#         In the second case, skiptile, just concat
#         """
#         batch_size_tensor = tf.shape(input_tensor)[0]
#         x_dim = tf.shape(input_tensor)[1]
#         y_dim = tf.shape(input_tensor)[2]
#         xx_ones = tf.ones([batch_size_tensor, x_dim], dtype=tf.int32)
#         xx_ones = tf.expand_dims(xx_ones, -1)
#         xx_range = tf.tile(tf.expand_dims(tf.range(y_dim), 0), [batch_size_tensor, 1])
#         xx_range = tf.expand_dims(xx_range, 1)
#         xx_channel = tf.matmul(xx_ones, xx_range)
#         xx_channel = tf.expand_dims(xx_channel, -1)
#         yy_ones = tf.ones([batch_size_tensor, y_dim], dtype=tf.int32)
#         yy_ones = tf.expand_dims(yy_ones, 1)
#         yy_range = tf.tile(tf.expand_dims(tf.range(x_dim), 0), [batch_size_tensor, 1])
#         yy_range = tf.expand_dims(yy_range, -1)
#         yy_channel = tf.matmul(yy_range, yy_ones)
#         yy_channel = tf.expand_dims(yy_channel, -1)
# 
#         x_dim = tf.cast(x_dim, tf.float32)
#         y_dim = tf.cast(y_dim, tf.float32)
# 
#         xx_channel = tf.cast(xx_channel, tf.float32) / (x_dim - 1)
#         yy_channel = tf.cast(yy_channel, tf.float32) / (y_dim - 1)
#         xx_channel = xx_channel * 2 - 1
#         yy_channel = yy_channel * 2 - 1
# 
#         ret = tf.concat([input_tensor, xx_channel, yy_channel], axis=-1)
# 
#         output_tensor = ret
# 
#         return output_tensor


# Path: models/illumination_integration_net/cnn.py
import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
from nn_utils.activations import from_hdr_activation, to_hdr_activation
from nn_utils.coord_conv import AddCoords

                output_units,
                1,
                padding="same",
                activation=to_hdr_activation,
            )
        )

        self.net = tf.keras.Sequential(net)

    @tf.function
    def call(self, embedding):
        reconstruction = self.net(embedding)

        return reconstruction


class CnnDiscriminator(tf.keras.Model):
    def __init__(
        self,
        discriminator_units: int = 32,
        activation="relu",
        img_height: int = 128,
        **kwargs,
    ):
        super(CnnDiscriminator, self).__init__(**kwargs)

        activation = tf.keras.layers.Activation(activation)

        net = [
            tf.keras.layers.InputLayer((img_height, img_height * 2, 3)),

            tf.keras.layers.Lambda(lambda x: from_hdr_activation(x)),