Added architectures

.gitignore

@@ -1 +1,2 @@
-venv
+venv
+__pycache__

api.py

@@ -8,7 +8,7 @@ import io
 import base64
 import logging
 from PIL import Image
-#from architectures import *
+from architectures import *
 
 # Logging configuration for observability and debugging
 logging.basicConfig(level=logging.INFO)

architectures/__init__.py

@@ -0,0 +1,7 @@
+from .average import Average
+from .cnn import CNNUpscaler
+from .reswish import Reswish
+from .espcn import ESPCN
+from .srgan import SRGAN
+from .srrn import SRRN
+

architectures/average.py

@@ -0,0 +1,50 @@
+import numpy as np
+import keras
+import tensorflow as tf
+from keras import ops
+from keras import layers
+
+@keras.saving.register_keras_serializable()
+class Upscaler(layers.Layer):
+    """Upscales images by superposing grids and averaging colors."""
+
+    def __init__(self, up_ratio: float, name="upscaler", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+
+    def call(self, inputs):
+        shape = tf.shape(inputs)
+        height = shape[1]
+        width = shape[2]
+
+        # We round the new width and height casting twice (tf is weird)
+        new_height = tf.cast(height, tf.float32) * self.up_ratio
+        new_width = tf.cast(width, tf.float32) * self.up_ratio
+        
+        new_height = tf.cast(new_height, tf.int32)
+        new_width = tf.cast(new_width, tf.int32)
+
+        # Resize
+        return tf.image.resize(inputs, [new_height, new_width], method='bilinear')
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio})
+        return config
+
+@keras.saving.register_keras_serializable()
+class Average(keras.Model):
+    """Defines a model which upscales images by averaging. Training does not modify its behavior"""
+
+    def __init__(self, up_ratio=2.0, name="average", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+        self.upscaler = Upscaler(up_ratio)
+
+    def call(self, inputs):
+        return self.upscaler(inputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio})
+        return config

architectures/cnn.py

@@ -0,0 +1,33 @@
+import keras
+import tensorflow as tf
+from keras import layers
+from .average import Upscaler
+
+@keras.saving.register_keras_serializable()
+class CNNUpscaler(keras.Model):
+    """Convolutional Neural Network for image upscaling. It's just a simple model: upscales first, then applies convolution corrections."""
+
+    def __init__(self, up_ratio: float, name="cnnupscaler", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+        
+        self.upscaler = Upscaler(up_ratio)
+        
+        self.conv1 = layers.Conv2D(64, (3, 3), activation="relu", padding="same")
+        self.conv2 = layers.Conv2D(32, (3, 3), activation="relu", padding="same")
+        self.conv3 = layers.Conv2D(3, (3, 3), padding="same")
+
+    def call(self, inputs):
+        # Upscale first
+        x_up = self.upscaler(inputs)
+
+        # Calculate the correction factor
+        x = self.conv1(x_up)
+        x = self.conv2(x)
+        correction = self.conv3(x)
+        return x_up + correction
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio})
+        return config

architectures/espcn.py

@@ -0,0 +1,38 @@
+import keras
+import tensorflow as tf
+from keras import layers
+
+@keras.saving.register_keras_serializable()
+class ESPCN(keras.Model):
+    """
+    Efficient Sub-Pixel Convolutional Neural Network for image upscaling.
+    Works in LR space, uses pixel shuffle at the end for fast and stable super-resolution.
+    """
+    
+    def __init__(self, up_ratio, name="espcn", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+
+        self.conv1= layers.Conv2D(64, 3, padding='same', activation='relu')
+        self.conv2 = layers.Conv2D(64, 3, padding='same', activation='relu')
+        self.conv3 = layers.Conv2D(32, 3, padding='same', activation='relu')
+        self.conv4 = layers.Conv2D(3 * (self.up_ratio ** 2), 3, padding='same')
+        
+        self.pixel_shuffle = layers.Lambda(lambda x: tf.nn.depth_to_space(x, block_size=self.up_ratio))
+
+    def call(self, inputs):
+        # Calculate the corrections
+        x = self.conv1(inputs)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        
+        # Upscale by pixel_suffle
+        x = self.pixel_shuffle(x)
+
+        return tf.clip_by_value(x, 0.0, 1.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio})
+        return config

architectures/reswish.py

@@ -0,0 +1,205 @@
+import tensorflow as tf
+import keras
+from keras import layers
+
+
+@keras.saving.register_keras_serializable()
+class ChannelAttention(layers.Layer):
+    """
+    Channel attention https://arxiv.org/pdf/1807.02758
+    """
+
+    def __init__(self, channels: int, reduction: int = 16, **kwargs):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.reduction = reduction
+
+        # Global Average Pooling
+        self.gap = layers.GlobalAveragePooling2D()
+
+        # Little bit of dense
+        self.dense1 = layers.Dense(channels // reduction, activation="relu")
+        self.dense2 = layers.Dense(channels, activation="sigmoid")
+
+        # Reshape
+        self.reshape = layers.Reshape((1, 1, channels))
+
+    def call(self, inputs):
+        x = self.gap(inputs)
+        x = self.dense1(x)
+        x = self.dense2(x)
+        x = self.reshape(x)
+
+        return inputs * x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"channels": self.channels, "reduction": self.reduction})
+        return config
+
+
+@keras.saving.register_keras_serializable()
+class ReswishLayer(layers.Layer):
+    """
+    A ReswishLayer is a residual block which aims to combine convolutions with trainable activation layers in order to produce complex models capable of learning different behaviors.
+    """
+
+    def __init__(self, filters: int, **kwargs):
+        """Generates the layers in a ReswishLayer.
+
+        Args:
+            filters (int): Amount of filters wanted after convoluting
+        """
+        super().__init__(**kwargs)
+
+        # Save stuff for config
+        self.filters = filters
+
+        # 2D Convolution
+        self.conv = layers.Conv2D(filters, (3, 3), padding="same")
+        # Swish it
+        self.sw1 = layers.Activation("swish")
+
+        # Conv again
+        self.conv2 = layers.Conv2D(filters, (3, 3), padding="same")
+        
+        # Channel attention
+        self.ca = ChannelAttention(filters, reduction=16)
+
+    def call(self, inputs):
+        res = inputs
+
+        x = self.conv(inputs)
+        x = self.sw1(x)
+        x = self.conv2(x)
+        x = self.ca(x)
+
+        return layers.add([res, x])
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"filters": self.filters})
+        return config
+
+
+@keras.saving.register_keras_serializable()
+class PixelShuffleUpscale(layers.Layer):
+    """
+    Upscales by rearranging additional channels into pixels
+    """
+
+    def __init__(self, shuffle_num: int, filters: int, **kwargs):
+        super().__init__(**kwargs)
+        # Save stuff for config
+        self.filters = filters
+        self.shuffle_num = shuffle_num
+
+        # Actual shuffle layers
+        self.conv = layers.Conv2D(
+            filters * shuffle_num * shuffle_num, (3, 3), padding="same"
+        )
+
+    def call(self, inputs):
+        x = self.conv(inputs)
+        x = tf.nn.depth_to_space(x, block_size=self.shuffle_num)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"filters": self.filters, "shuffle_num": self.shuffle_num})
+        return config
+
+
+@keras.saving.register_keras_serializable()
+class ColorConstraint(layers.Layer):
+    """
+    Calculates the color consistency error and adds it to the model's loss
+    without modifying the image pixels directly.
+    """
+
+    def __init__(self, up_ratio: int, weight: float = 0.5, **kwargs):
+        super().__init__(**kwargs)
+        self.up_ratio = up_ratio
+        self.weight = weight
+
+    def call(self, inputs):
+        # inputs must be [gen, og]
+        gen, og = inputs
+
+        # Downscale to compare
+        gen_downscaled = tf.nn.avg_pool2d(
+            gen, ksize=self.up_ratio, strides=self.up_ratio, padding="VALID"
+        )
+
+        # Get error
+        color_error = tf.reduce_mean(tf.abs(og - gen_downscaled))
+
+        # Add loss
+        self.add_loss(color_error * self.weight)
+
+        return gen
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0]
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio, "weight": self.weight})
+        return config
+
+
+@keras.saving.register_keras_serializable()
+class Reswish(keras.Model):
+    """Defines a model using ReswishLayer architecture."""
+
+    def __init__(
+        self,
+        filters: int = 64,
+        blocks: int = 4,
+        up_ratio: int = 5,
+        name="Reswish",
+        **kwargs
+    ):
+        # Save stuff for config
+        self.filters = filters
+        self.blocks = blocks
+        self.up_ratio = up_ratio
+
+        # Declare input (RGB images)
+        inputs = layers.Input(shape=(None, None, 3))
+
+        # Expand to #filters
+        x = layers.Conv2D(filters, (3, 3), padding="same")(inputs)
+
+        # Keep residual
+        res = x
+
+        # Add ReswishLayers
+        for _ in range(blocks):
+            x = ReswishLayer(filters)(x)
+
+        # Extra conv to keep it on its toes
+        x = layers.Conv2D(filters, (3, 3), padding="same")(x)
+
+        # Add res back
+        x = layers.add([x, res])
+
+        # Upscale
+        x = PixelShuffleUpscale(filters=filters, shuffle_num=up_ratio)(x)
+
+        # Collapse to RGB
+        x = layers.Conv2D(3, (3, 3), padding="same", activation=None)(x)
+
+        # Correct
+        outputs = ColorConstraint(up_ratio)([x, inputs])
+
+        # Create whole model
+        super().__init__(inputs, outputs, name=name, **kwargs)
+
+    def get_config(self):
+        # This allows model.save() to work
+        return {
+            "filters": self.filters,
+            "blocks": self.blocks,
+            "up_ratio": self.up_ratio,
+        }

architectures/srgan.py

@@ -0,0 +1,141 @@
+import tensorflow as tf
+import keras
+from keras import layers
+
+# import the generator from other modules (or create a custom one)
+from .srrn import SRRN
+
+@keras.saving.register_keras_serializable()
+class DiscBlock(layers.Layer):
+    """Basic discriminator block: Conv -> BatchNorm -> LeakyReLu"""
+    def __init__(self, filters, stride=1, name="DiscBlock"):
+        super().__init__(name=name)
+        self.conv = layers.Conv2D(filters, 3, strides=stride, padding="same")
+        self.bn = layers.BatchNormalization()
+        self.act = layers.LeakyReLU(0.2)
+
+    def call(self, x, training=False):
+        x = self.conv(x)
+        x = self.bn(x, training=training)
+        return self.act(x)
+
+@keras.saving.register_keras_serializable()
+class Discriminator(keras.Model):
+    """Discriminator for SRGAN"""
+    def __init__(self, name="discriminator", **kwargs):
+        super().__init__(name=name, **kwargs)
+        
+        # Discriminator blocks
+        self.blocks = [
+            DiscBlock(64, stride=1),
+            DiscBlock(64, stride=2),
+
+            DiscBlock(128, stride=1),
+            DiscBlock(128, stride=2),
+
+            DiscBlock(256, stride=1),
+            DiscBlock(256, stride=2),
+
+            DiscBlock(512, stride=1),
+            DiscBlock(512, stride=2)
+        ]
+
+        # Final conv
+        self.final_conv = layers.Conv2D(1, 3, padding="same")
+
+    def call(self, inputs, training=False):
+        x = inputs
+        for block in self.blocks:
+            x = block(x, training=training)
+        x = self.final_conv(x)
+        return x
+    
+    def get_config(self):
+        config = super().get_config()
+        return config
+
+@keras.saving.register_keras_serializable()
+class SRGAN(keras.Model):
+    """Super-Resolution Generator Adversarial Network. Includes the generator + discriminator + custom train step"""
+    def __init__(self, up_ratio, filters=64, num_blocks=8, lambda_adv=1e-3, name="srgan", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+
+        # Generator (SRResNet)
+        self.generator = SRRN(up_ratio=up_ratio, filters=filters, num_blocks=num_blocks,)
+
+        # Discriminator
+        self.discriminator = Discriminator()
+
+        # Loss weights
+        self.lambda_adv = lambda_adv
+
+        # Loss functions
+        self.pixel_loss_fn = tf.keras.losses.MeanAbsoluteError()
+
+        # Optimizers (they are assigned in the compile fase)
+        self.gen_optimizer = None
+        self.disc_optimizer = None
+
+    def discriminator_loss(self, real_logits, fake_logits):
+        # Uses Least Squares instead of Binary Cross-Entropy
+        real_loss = tf.reduce_mean((real_logits - 1.0) ** 2)
+        fake_loss = tf.reduce_mean((fake_logits - 0.0) ** 2)
+        return real_loss + fake_loss
+
+    def _generator_adversarial_loss(self, fake_logits):
+        return tf.reduce_mean((fake_logits - 1.0) ** 2)
+
+    def generator_total_loss(self, sr, hr, fake_logits):
+        pixel_loss = self.pixel_loss_fn(hr, sr)
+        adv_loss = self._generator_adversarial_loss(fake_logits)
+        return pixel_loss + self.lambda_adv * adv_loss
+
+    def compile(self, gen_optimizer, disc_optimizer, **kwargs):
+        super().compile(**kwargs)
+
+        self.gen_optimizer = gen_optimizer
+        self.disc_optimizer = disc_optimizer
+
+    def call(self, inputs):
+        return self.generator(inputs)
+    
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio, "lambda_adv": self.lambda_adv})
+        return config
+
+    @tf.function
+    def train_step(self, data):
+        lr, hr = data
+
+        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
+            # Generate SR image
+            sr = self.generator(lr, training=True)
+
+            # Discriminator predictions
+            real_logits = self.discriminator(hr, training=True)
+            fake_logits = self.discriminator(sr, training=True)
+
+            # Losses
+            disc_loss = self.discriminator_loss(real_logits, fake_logits)
+            gen_loss = self.generator_total_loss(sr, hr, fake_logits)
+
+        # Gradients
+        gen_grads = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
+        disc_grads = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
+
+        # Apply gradients
+        self.gen_optimizer.apply_gradients(zip(gen_grads, self.generator.trainable_variables))
+        self.disc_optimizer.apply_gradients(zip(disc_grads, self.discriminator.trainable_variables))
+
+        # PSNR y SSIM metri 
+        psnr_val = tf.image.psnr(hr, sr, max_val=1.0)
+        ssim_val = tf.image.ssim(hr, sr, max_val=1.0)
+
+        return {
+            "gen_loss": gen_loss,
+            "disc_loss": disc_loss,
+            "psnr": tf.reduce_mean(psnr_val),
+            "ssim": tf.reduce_mean(ssim_val)
+            }

architectures/srrn.py

@@ -0,0 +1,69 @@
+import keras
+import tensorflow as tf
+from keras import layers
+
+@keras.saving.register_keras_serializable()
+class ResidualBlock(layers.Layer):
+    """Basic residual block: Conv -> ReLU"""
+    def __init__(self, filters, **kwargs):
+        super().__init__(**kwargs)
+        self.conv1 = layers.Conv2D(filters, 3, padding='same')
+        self.relu = layers.ReLU()
+        self.conv2 = layers.Conv2D(filters, 3, padding='same')
+
+    def call(self, x, training=False):
+        residual = x
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        return x + residual
+
+@keras.saving.register_keras_serializable()
+class SRRN(keras.Model):
+    """
+    Super-Resolution Residual Network. A more sofisticated post-upscaler Network.
+    Uses residual blocks to refine features in LR space, then performs upsampling via PixelShuffle.
+    """
+
+    def __init__(self, up_ratio, filters=64, num_blocks=8, name="SRRN", **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.up_ratio = up_ratio
+
+        # Intial conv
+        self.conv_in = layers.Conv2D(filters, 9, padding='same')
+        self.relu = layers.ReLU()
+
+        # Residual blocks
+        self.res_blocks = [ResidualBlock(filters) for _ in range(num_blocks)]
+
+        # Post residual
+        self.conv_post_res = layers.Conv2D(filters, 3, padding='same')
+
+        # Upsampling
+        self.upsample = layers.Conv2D(filters * (up_ratio**2), 3, padding='same')
+        self.pixel_shuffle = layers.Lambda(lambda x: tf.nn.depth_to_space(x, up_ratio))
+
+        # Final conv
+        self.conv_final = layers.Conv2D(3, 9, padding='same', activation='sigmoid')
+
+    def call(self, x, training=False):
+        x = self.conv_in(x)
+        x = self.relu(x)
+
+        # Residual blocks
+        res = x
+        for block in self.res_blocks:
+            x = block(x, training=training)
+        x = self.conv_post_res(x)
+        x = layers.add([res, x])
+
+        # Upsampling
+        x = self.upsample(x)
+        x = self.pixel_shuffle(x)
+
+        return self.conv_final(x)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"up_ratio": self.up_ratio})
+        return config