Initial commit

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/github_tide.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="tf2" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="GOOGLE" />
+    <option name="myDocStringFormat" value="Google" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,10 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="DuplicatedCode" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <Languages>
+        <language minSize="66" name="Python" />
+      </Languages>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Black">
+    <option name="sdkName" value="Python 3.12" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="tf2" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/github_tide.iml" filepath="$PROJECT_DIR$/.idea/github_tide.iml" />
+    </modules>
+  </component>
+</project>

compare_ckpts.py

@@ -0,0 +1,192 @@
+import sys
+import glob
+import numpy as np
+import pandas as pd
+import importlib
+import tensorflow as tf
+
+from PIL import Image
+from re import split, compile
+from tensorflow.keras.applications.inception_v3 import preprocess_input
+
+import fid_kid
+
+#TODO : uncomment & import appropriate paths if msgastrovae_smc.py (TIDE model) & my_convnext.py (TIDE-2 model)
+#       are not in the same folder with this script (for importlib modules below)
+# sys.path.append(f'{tide_path}')
+# sys.path.append(f'{tide2_path}')
+
+
+def list_saved_models(results_dir):
+    models_found = glob.glob("{}/weights/vae_checkpoints/*.index".format(results_dir))    # checkpoints
+    models_found.extend(glob.glob("{}/weights/*.h5".format(results_dir)))                 # models
+    models_found.sort(key=lambda l: [int(s) if s.isdigit() else s.lower() for s in split(compile(r'(\d+)'), l)])
+    # print(models_found)
+    return models_found
+
+
+def get_real_kid_filenames(label):
+    kid_dir = '/mnt/storage/shared/ckaitanidis/datasets/kid/kid-dataset-2/'
+    files = []
+    if label in ['inflammatory', 'polypoid', 'vascular']:
+        files = glob.glob('{}/{}/*.png'.format(kid_dir, label))
+    elif label == 'normaleso':
+        files = glob.glob('{}/normal-esophagus'.format(kid_dir))
+    elif label == 'normalstom':
+        files = glob.glob('{}/normal-stomach'.format(kid_dir))
+    elif label == 'normalcolon':
+        files = glob.glob('{}/normal-colon'.format(kid_dir))
+    elif label == 'normalsb':
+        files = glob.glob('{}/normal-small-bowel'.format(kid_dir))
+    print('Real images found: {}'.format(len(files)))
+    files = sorted(files)
+    return files
+
+
+def init_vae_model(model_name, latent_dim, input_shape):
+    if model_name == 'tide':
+        vae = importlib.import_module("msgastrovae_smc")
+        vae_model = vae.VAE(vae.create_encoder(latent_dim=latent_dim,input_shape=input_shape),
+                            vae.create_decoder(latent_dim=latent_dim))
+        vae_model.build(input_shape=[(None,) + input_shape])
+        return vae_model
+    elif model_name == 'tide2':
+        vae = importlib.import_module("my_convnext")
+        vae_model = vae.VAE(vae.create_encoder_tiny(latent_dim=latent_dim, input_shape=input_shape),
+                            vae.create_decoder_tiny(latent_dim=latent_dim))
+        vae_model.build(input_shape=[(None,) + input_shape])
+        return vae_model
+
+
+def load_weights(vae, weights_path):
+    print("Loading weights from {}".format(weights_path))
+    if "ckpt-" in weights_path:
+        weights_path = weights_path.split(".index")[0]
+        ckpt = tf.train.Checkpoint(vae=vae)
+        status = ckpt.restore(weights_path).expect_partial()
+        return vae
+    if ".h5" in weights_path:
+        vae.load_weights(weights_path, by_name=True)
+        return vae
+
+
+def debug_weights_loading(vae):
+    decoder_weights = vae.decoder.get_weights()
+    print("Decoder layer 0 weights shape:", decoder_weights[0].shape)
+    print("Decoder layer 0 weights sample:", decoder_weights[0].flatten()[:5])
+    # 'Decoder layer 0 weights sample: [-0.01202846 -0.02691004  0.00642165 -0.02967337 -0.03743371]' # mine always same
+
+
+def get_noise_seeded(noise_shape):
+    np.random.seed(0)
+    random_z = np.random.normal(0, 1, noise_shape)
+    return random_z
+
+
+def decode_noise(trained_vae, noise, return_list=False):
+    print("Generating fake images ...")
+    pred = trained_vae.decoder.predict(noise, batch_size=1)
+    # print(type(pred), pred.shape, pred.dtype, pred.min(), pred.max())
+    pred *= 255.0   # for tf.preprocess_input requires [0, 255]
+    # print(type(pred), pred.shape, pred.dtype, pred.min(), pred.max())
+    if return_list:
+        return [img for img in pred]
+    return pred
+
+
+def visualize_debug(image, name='output.png'):
+    image = ((image + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
+    Image.fromarray(image).save(name)
+
+
+def kid_dataset_center_crop(img, crop_size=(320, 320)):
+    if not isinstance(img, np.ndarray):
+        img = np.array(img)
+
+    h, w, _ = img.shape
+    ch, cw = crop_size
+
+    top = (h - ch) // 2
+    left = (w - cw) // 2
+
+    return img[top:top + ch, left:left + cw]
+
+
+if __name__ == "__main__":
+
+    # Change these two to run across all models
+    model_name = 'tide2'             # 'tide'  'tide2'
+    label = "inflammatory"           # 'inflammatory' 'vascular' 'polypoid' 'normaleso' 'normalcolon' 'normalsb' 'normalstom'
+
+    results_dir = "/mnt/storage/shared/ckaitanidis/"
+    if model_name == 'tide':
+        results_dir += 'kid_latent6_sr96_50000ep_{}'.format(label)
+    elif model_name == 'tide2':
+        results_dir += 'kid_latent8_tide2_sr96_50000ep_{}'.format(label)
+    print(results_dir)
+
+    # Params auto
+    latent_dim = 6 if model_name == 'tide' else 8
+    input_shape = (96, 96, 3) if model_name == 'tide' else (256, 256, 3)
+    crop_dim = (320, 320)
+
+    trained_weights = list_saved_models(results_dir)
+    real_filenames = get_real_kid_filenames(label)
+    real_images = fid_kid.get_images_inception(real_filenames, crop_dim=crop_dim)  # this returns np.array, float32, [-1, 1], (batch, 299, 299, 3)
+    # print(type(real_images), real_images.shape, real_images.dtype, real_images.min(), real_images.max())
+    visualize_debug(real_images[0], name='output1.png')
+
+    vae = init_vae_model(model_name, latent_dim, input_shape)
+    noise_vector = get_noise_seeded((len(real_filenames), latent_dim))
+
+
+    results = []
+
+    # Ignore these - my weights for debug
+    # trained_weights = ['/mnt/storage/pgatoula-private/codes/tide-panagiota/results_kid/kid_inflammatory_latent6/weights/vae_checkpoints/ckpt-4500.index']
+    # trained_weights = ['/mnt/storage/pgatoula-private/general-results/convnext/kid_inflammatory_latent8_lbfcn_sr96/weights/vae_checkpoints/ckpt-1400.index',
+    #                    '/mnt/storage/pgatoula-private/general-results/convnext/kid_inflammatory_latent8_lbfcn_sr96/weights/vae_checkpoints/ckpt-1600.index']
+
+    for weights in trained_weights:
+        # Load weights
+        vae = load_weights(vae, weights)
+        vae.trainable = False
+        # try:
+        #     debug_weights_loading(vae)
+        # except Exception as e:
+        #     print(f"Skipping {weights} due to load failure: {e}")
+        #     continue
+
+        # Generate Fakes
+        fake_images = decode_noise(vae, noise_vector, return_list=True)
+        fake_images = preprocess_input(fake_images)
+        # print(type(fake_images), fake_images.shape, fake_images.dtype, fake_images.min(), fake_images.max())
+        fake_images = tf.image.resize(fake_images, size=(299, 299), method='bicubic').numpy()
+        # print(type(fake_images), fake_images.shape, fake_images.dtype, fake_images.min(), fake_images.max())
+        visualize_debug(fake_images[0], name='output2.png')
+
+        # Calculate metrics
+        fid_score = fid_kid.calculate_fid(real_images, fake_images)
+        kid_mean, kid_std = fid_kid.calculate_kid(real_images, fake_images)
+
+        fid_score = round(fid_score, 4)
+        kid_mean = round(kid_mean, 4)
+        kid_std = round(kid_std, 4)
+
+        print("{}: FID={} KID={} ± {}".format(weights, fid_score, kid_mean, kid_std))
+
+        results.append({'weights': weights,
+                        'fid': fid_score,
+                        'kid_mean': kid_mean,
+                        'kid_std': kid_std,
+                        })
+
+    # Save in xlxs
+    df = pd.DataFrame(results)
+    excel_path = f"ckpt_metrics_{model_name}.xlsx"
+
+    # TODO: pip install XlsxWriter if not installed
+    with pd.ExcelWriter(excel_path, engine='xlsxwriter') as writer:
+        df.to_excel(writer, sheet_name=label, index=False)
+
+    print(f"Results saved to: {excel_path}")

generate_images.py

@@ -0,0 +1,30 @@
+import os
+
+from argparse import ArgumentParser
+from utils.inference_utils import init_vae_model, load_weights, get_noise_seeded, decode_noise, save_images
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--model_name", required=True, type=str, choices=['tide', 'tidev2'], help='VAE model')
+    parser.add_argument("--weights_path", required=True, type=str, help='Path to restore trained weights')
+    parser.add_argument("--latent_dim", default=8, type=int, help='Dimensionality of latent space')
+    parser.add_argument("--save_dir", default="./fake_images", type=str, help='Path to save synthetic images')
+    parser.add_argument("--num_of_images", default=10, type=int, help='Number of images to generate')
+    args = parser.parse_args()
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    if not os.path.exists(args.weights_path):
+        print("Not a valid path")
+
+    vae = init_vae_model(args.model_name, args.latent_dim)
+    noise_vector = get_noise_seeded((args.num_of_images, args.latent_dim))
+
+    # Load weights
+    vae = load_weights(vae, args.weights_path)
+    vae.trainable = False
+
+    # Generate & Save images
+    fake_images = decode_noise(vae, noise_vector, return_list=True)
+    save_images(args.save_dir, fake_images)

model/convnext_modules.py

@@ -0,0 +1,111 @@
+import tensorflow as tf
+import tensorflow.keras.layers as layers
+
+from tensorflow.keras import backend
+
+
+class LayerScale(layers.Layer):
+    def __init__(self, init_values, projection_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.init_values = init_values
+        self.projection_dim = projection_dim
+
+    def build(self, input_shape):
+        self.gamma = tf.Variable(self.init_values * tf.ones((self.projection_dim,)))
+
+    def call(self, x):
+        return x * self.gamma
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "init_values": self.init_values,
+                "projection_dim": self.projection_dim,
+            }
+        )
+        return config
+
+
+class StochasticDepth(layers.Layer):
+    def __init__(self, drop_path_rate, **kwargs):
+        super().__init__(**kwargs)
+        self.drop_path_rate = drop_path_rate
+
+    def call(self, x, training=None):
+        if training:
+            keep_prob = 1 - self.drop_path_rate
+            shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
+            random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
+            random_tensor = tf.floor(random_tensor)
+            return (x / keep_prob) * random_tensor
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"drop_path_rate": self.drop_path_rate})
+        return config
+
+
+class ConvNeXtBlock(layers.Layer):
+    def __init__(self, projection_dim, drop_path_rate=0.0, layer_scale_init_value=1e-6, name_prefix=None):
+        super().__init__(name=name_prefix or f"prestem{backend.get_uid('prestem')}")
+        self.depthwise_conv = layers.Conv2D(
+            filters=projection_dim, kernel_size=7, padding="same", groups=projection_dim,
+            name=self.name + "_depthwise_conv"
+        )
+        self.pointwise_conv1 = layers.Dense(4 * projection_dim, name=self.name + "_pointwise_conv_1")
+        self.act = layers.Activation("gelu", name=self.name + "_gelu")
+        self.pointwise_conv2 = layers.Dense(projection_dim, name=self.name + "_pointwise_conv_2")
+        self.layer_scale = LayerScale(layer_scale_init_value, projection_dim, name=self.name + "_layer_scale") \
+            if layer_scale_init_value is not None else None
+        self.stochastic_depth = StochasticDepth(drop_path_rate, name=self.name + "_stochastic_depth") \
+            if drop_path_rate else layers.Activation("linear", name=self.name + "_identity")
+
+    def call(self, inputs, training=False):
+        x = self.depthwise_conv(inputs)
+        x = self.pointwise_conv1(x)
+        x = self.act(x)
+        x = self.pointwise_conv2(x)
+        if self.layer_scale:
+            x = self.layer_scale(x)
+        x = self.stochastic_depth(x, training=training)
+        return inputs + x
+
+
+class ConvNeXtBlockTransposed(layers.Layer):
+    def __init__(self, projection_dim, drop_path_rate=0.0, layer_scale_init_value=1e-6, name_prefix=None):
+        super().__init__(name=name_prefix or f"poststem{backend.get_uid('poststem')}")
+        self.projection_dim = projection_dim
+        self.drop_path_rate = drop_path_rate
+        self.layer_scale_init_value = layer_scale_init_value
+
+        self.depthwise_conv_trans = layers.Conv2DTranspose(
+            filters=projection_dim, kernel_size=7, padding="same",
+            groups=projection_dim, name=self.name + "_depthwise_conv_trans"
+        )
+        self.pointwise_conv1 = layers.Dense(4 * projection_dim, name=self.name + "_pointwise_conv_1")
+        self.act = layers.Activation("gelu", name=self.name + "_gelu")
+        self.pointwise_conv2 = layers.Dense(projection_dim, name=self.name + "_pointwise_conv_2")
+
+        if layer_scale_init_value is not None:
+            self.layer_scale = LayerScale(layer_scale_init_value, projection_dim, name=self.name + "_layer_scale")
+        else:
+            self.layer_scale = None
+
+        if drop_path_rate:
+            self.stochastic_depth = StochasticDepth(drop_path_rate, name=self.name + "_stochastic_depth")
+        else:
+            self.stochastic_depth = layers.Activation("linear", name=self.name + "_identity")
+
+    def call(self, inputs, training=False):
+        x = self.depthwise_conv_trans(inputs)
+        x = self.pointwise_conv1(x)
+        x = self.act(x)
+        x = self.pointwise_conv2(x)
+        if self.layer_scale:
+            x = self.layer_scale(x)
+        x = self.stochastic_depth(x, training=training)
+        return inputs + x
+
+

model/tidev2.py

@@ -0,0 +1,161 @@
+import numpy as np
+import tensorflow.keras.layers as layers
+
+from tensorflow.keras import Model
+from tensorflow.keras import Sequential
+
+from model.tidev2_utils import TopLayer, Sampling
+from model.convnext_modules import ConvNeXtBlock, ConvNeXtBlockTransposed
+
+
+class ConvNeXtEncoderTiny(Model):
+    def __init__(self,
+                 depths=[3, 3, 9, 3],
+                 projection_dims=[96, 192, 384, 768],
+                 drop_path_rate=0.0,
+                 layer_scale_init_value=1e-6,
+                 model_name="convnext",
+                 latent_dim=None):
+        super().__init__(name=model_name)
+        self.latent_dim = latent_dim
+        self.depths = depths
+        self.projection_dims = projection_dims
+
+        # Stem
+        self.stem = Sequential([
+            layers.Conv2D(projection_dims[0], kernel_size=4, strides=4, name=model_name + "_stem_conv"),
+        ], name=model_name + "_stem")
+
+        # Downsampling layers
+        self.downsample_layers = [self.stem]
+        for i in range(3):
+            self.downsample_layers.append(
+                Sequential([
+                    layers.Conv2D(projection_dims[i + 1], kernel_size=2, strides=2,
+                                  name=model_name + f"_downsampling_conv_{i}")
+                ], name=model_name + f"_downsampling_block_{i}")
+            )
+
+        # Drop rates for stochastic depth
+        self.depth_drop_rates = np.linspace(0.0, drop_path_rate, sum(depths)).astype(float)
+
+        # ConvNeXt stages
+        self.stages = []
+        cur = 0
+        for i in range(4):
+            stage_blocks = []
+            for j in range(depths[i]):
+                stage_blocks.append(
+                    ConvNeXtBlock(projection_dim=projection_dims[i],
+                                  drop_path_rate=self.depth_drop_rates[cur + j],
+                                  layer_scale_init_value=layer_scale_init_value,
+                                  name_prefix=model_name + f"_stage_{i}_block_{j}")
+                )
+            self.stages.append(stage_blocks)
+            cur += depths[i]
+
+        # Latent projection if requested
+        if latent_dim is not None:
+            self.flatten = layers.Flatten()
+            self.dense_proj = layers.Dense(256, activation="relu", name="dense_proj")
+            self.z_mean = layers.Dense(latent_dim, name="z_mean")
+            self.z_log_var = layers.Dense(latent_dim, name="z_log_var")
+            self.sampling = Sampling()
+
+    def call(self, inputs, training=False):
+        x = inputs
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            for block in self.stages[i]:
+                x = block(x, training=training)
+
+        if self.latent_dim is None:
+            return x
+
+        x = self.flatten(x)
+        x = self.dense_proj(x)
+        z_mean = self.z_mean(x)
+        z_log_var = self.z_log_var(x)
+        z = self.sampling([z_mean, z_log_var])
+        return [z, z_mean, z_log_var]
+
+
+class ConvNeXtDecoderTiny(Model):
+    def __init__(self,
+                 depths=[3, 9, 3, 3],
+                 projection_dims=[768, 384, 192, 96],
+                 drop_path_rate=0.0,
+                 layer_scale_init_value=1e-6,
+                 model_name="convnext",
+                 latent_dim=None):
+        super().__init__(name=model_name)
+
+        if latent_dim is None:
+            raise ValueError("latent_dim must be specified for decoder")
+
+        # Intro layer (dense + reshape)
+        self.intro = Sequential([
+            layers.Dense(10 * 10 * projection_dims[0], activation="relu"),
+            layers.Reshape((10, 10, projection_dims[0]))
+        ], name=model_name + "_intro")
+
+        # Upsampling layers
+        self.upsample_layers = [self.intro]
+        for i in range(3):
+            self.upsample_layers.append(
+                Sequential([
+                    layers.Conv2DTranspose(projection_dims[i + 1], kernel_size=2, strides=2,
+                                           name=model_name + f"_upsampling_conv_{i}")
+                ], name=model_name + f"_upsampling_block_{i}")
+            )
+
+        # Drop rates for stochastic depth
+        self.depth_drop_rates = np.linspace(0.0, drop_path_rate, sum(depths)).astype(float)
+
+        # ConvNeXt transpose stages
+        self.stages = []
+        cur = 0
+        for i in range(4):
+            stage_blocks = []
+            for j in range(depths[i]):
+                stage_blocks.append(
+                    ConvNeXtBlockTransposed(projection_dim=projection_dims[i],
+                                            drop_path_rate=self.depth_drop_rates[cur + j],
+                                            layer_scale_init_value=layer_scale_init_value,
+                                            name_prefix=model_name + f"_stage_{i}_block_{j}")
+                )
+            self.stages.append(stage_blocks)
+            cur += depths[i]
+
+        # Top layer
+        self.top = Sequential([
+            layers.Conv2DTranspose(projection_dims[3], kernel_size=4, strides=4, name=model_name + "_top_conv")
+        ], name=model_name + "_top")
+
+        self.top_layer = TopLayer(filters=96)
+        self.pred_layer = layers.Conv2DTranspose(3, kernel_size=1, activation="sigmoid",
+                                                 padding="same", name="pred_layer")
+
+    def call(self, inputs, training=False):
+        x = inputs
+        for i in range(4):
+            x = self.upsample_layers[i](x)
+            for block in self.stages[i]:
+                x = block(x, training=training)
+        x = self.top(x)
+        x = self.top_layer(x)
+        return self.pred_layer(x)
+
+
+if __name__ == "__main__":
+    # Encoder
+    encoder = ConvNeXtEncoderTiny(latent_dim=8)
+    encoder.build((None, 320, 320, 3))
+    encoder.summary()
+
+    # Decoder
+    decoder = ConvNeXtDecoderTiny(latent_dim=8)
+    decoder.build((None, 8))
+    decoder.summary()
+
+

model/tidev2_utils.py

@@ -0,0 +1,56 @@
+import tensorflow as tf
+import tensorflow.keras.layers as layers
+
+
+class TopLayer(layers.Layer):
+    def __init__(self, filters):
+        super().__init__()
+        self.filters = filters
+
+        self.conv_1x1 = layers.Conv2D(self.filters, (1, 1), activation='relu', strides=1, padding="same",
+                                      name="_top_layer")
+        self.conv_2x2 = layers.Conv2D(self.filters//3, (2, 2), activation='relu', strides=1, padding="same",
+                                      name="_top_layer")
+        self.conv_4x4 = layers.Conv2D(self.filters//3, (4, 4), activation='relu', strides=1, padding="same",
+                                      name="_top_layer")
+        self.conv_8x8 = layers.Conv2D(self.filters//3, (8, 8), activation='relu', strides=1, padding="same",
+                                      name="_top_layer")
+
+        self.concat = layers.Concatenate(axis=-1)
+        self.point_wise_conv = layers.Conv2D(self.filters, (1, 1), 1, activation=None, use_bias=False,
+                                             padding='same', name="_top_layer")
+        self.feat_fusion = layers.Conv2D(self.filters, (1, 1), 1, activation=None, use_bias=False,
+                                         padding='same', name="_top_layer")
+
+        self.addition = layers.Add()
+        self.gelu = layers.Activation('gelu')
+        self.final_conv = layers.Conv2D(self.filters, (1, 1),  activation='relu', strides=1, padding="same",
+                                        name="_top_layer")
+
+    def call(self, inputs, training=False):
+        x = self.conv_1x1(inputs, training=training)
+
+        feats_2x2 = self.conv_2x2(x, training=training)
+        feats_4x4 = self.conv_4x4(x, training=training)
+        feats_8x8 = self.conv_8x8(x, training=training)
+
+        concatenated = self.concat([feats_2x2, feats_4x4, feats_8x8])
+        concatenated = self.point_wise_conv(concatenated)
+
+        concatenated = self.feat_fusion(concatenated)
+        x = self.addition([inputs, concatenated])
+        x = self.gelu(x)
+        x = self.final_conv(x)
+        return x
+
+
+class Sampling(layers.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, inputs):
+        z_mean, z_log_var = inputs
+        batch = tf.shape(z_mean)[0]
+        dim = tf.shape(z_mean)[1]
+        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
+        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

model/vae.py

@@ -0,0 +1,61 @@
+import tensorflow as tf
+from tensorflow.keras.models import Model
+
+
+class VAE(Model):
+    def __init__(self, encoder, decoder, **kwargs):
+        super(VAE, self).__init__(**kwargs)
+        self.encoder = encoder
+        self.decoder = decoder
+
+        # Loss Trackers
+        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
+        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(name="reconstruction_loss")
+        self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")
+
+    @property
+    def metrics(self):
+        return [
+            self.total_loss_tracker,
+            self.reconstruction_loss_tracker,
+            self.kl_loss_tracker,
+        ]
+
+    @tf.function()
+    def call(self, x):
+        z, z_mean, z_log_var, = self.encoder(x)
+        reconstruction = self.decoder(z)
+        return reconstruction
+
+    def full_summary(self):
+        for layer in self.layers:
+            print(layer.summary())
+
+    @tf.function()
+    def train_step(self, x):
+        with tf.GradientTape() as tape:
+            z, z_mean, z_log_var, = self.encoder(x)
+            reconstruction = self.decoder(z)
+
+            reconstruction_loss = tf.reduce_mean(
+                tf.reduce_sum(
+                    tf.keras.losses.binary_crossentropy(x, reconstruction), axis=(1, 2)
+                )
+            )
+            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
+            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
+            if tf.math.is_nan(kl_loss) or tf.math.is_inf(kl_loss):
+                kl_loss = tf.float32.max
+            total_loss = reconstruction_loss + kl_loss
+
+        grads = tape.gradient(total_loss, self.trainable_weights)
+        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
+        self.total_loss_tracker.update_state(total_loss)
+        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
+        self.kl_loss_tracker.update_state(kl_loss)
+
+        return {
+            "loss": self.total_loss_tracker.result(),
+            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
+            "kl_loss": self.kl_loss_tracker.result(),
+        }

train.py

@@ -0,0 +1,87 @@
+import os
+import tensorflow as tf
+
+from json import dump
+from argparse import ArgumentParser
+
+
+from model import tidev2
+from model.vae import VAE
+from utils.callbacks import VisualizeCallback, CheckpointCallback
+from utils.dataloader import list_filenames, Dataset
+from utils.plots import visualize_from_latent_space
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument("--model_name", required=True, type=str, choices=['tide', 'tidev2'], help='VAE model')
+    parser.add_argument("--output_path", default='./results/', type=str, help='Path to store the results')
+    # VAE model
+    parser.add_argument("--input_shape", default=(320, 320, 3), type=tuple, help='Image shape for training')
+    parser.add_argument("--dim_latent", default=8, type=int, help='Dimensionality of latent space')
+    # Training
+    parser.add_argument("--epochs", default=5000, type=int, help='Number of training epochs')
+    parser.add_argument("--batch_size", default=4, type=int, help='Number of training batch size')
+    parser.add_argument("--learning_rate",  default=0.0002, type=float, help='Learning rate')
+    parser.add_argument("--ckpt_interval", default=200, type=int, help='Epoch interval for saving checkpoints')
+    parser.add_argument("--visualization_interval", default=25, type=int, help='Epoch interval for visualizing results')
+    # Data
+    parser.add_argument("--datadir", default='./kid/inflammatory', type=str, help='Folder with images for training')
+    parser.add_argument("--files_ext", default='png', type=str, help='Extension of training files')
+    parser.add_argument("--files_prefix", default=None, type=str,
+                        help='Prefix of training files. Ignore if datadir contains only the appropriate files')
+    parser.add_argument("--crop_dim", default=None, type=tuple,
+                        help='Dimensions for cropping images. Ignore if images are already cropped')
+    args = parser.parse_args()
+
+    # Create folders & Save training config
+    os.makedirs(args.output_path, exist_ok=True)
+    log_dir = os.path.join(args.output_path, 'logs')
+    ckpt_dir = os.path.join(args.output_path, 'checkpoints')
+    visualize_dir = os.path.join(args.output_path, 'visualize')
+
+    os.makedirs(log_dir, exist_ok=True)
+    os.makedirs(ckpt_dir, exist_ok=True)
+    os.makedirs(visualize_dir, exist_ok=True)
+
+    with open(os.path.join(args.output_path, "training_config.json"), 'w') as fp:
+        dump(vars(args), fp)
+
+    # Setup training data
+    filenames = list_filenames(data_path=args.datadir,
+                               img_extension=args.files_ext,
+                               filename_prefix=args.files_prefix)
+    images = Dataset(filenames,
+                     batch_size=args.batch_size,
+                     crop_dim=args.crop_dim,
+                     resize_dim=args.input_shape[:2],)
+
+    # Create Model
+    if args.model_name == 'tidev2':
+        vae = VAE(tidev2.ConvNeXtEncoderTiny(latent_dim=args.dim_latent),
+                  tidev2.ConvNeXtDecoderTiny(latent_dim=args.dim_latent)
+                  )
+        vae.compile(optimizer=tf.keras.optimizers.Adam(args.learning_rate))
+
+    # Training
+    callbacks = [VisualizeCallback(args.visualization_interval, lambda model,  epoch: visualize_from_latent_space(
+                                                                                latent_dim=args.dim_latent,
+                                                                                input_shape=args.input_shape,
+                                                                                vae=model,
+                                                                                output_path=visualize_dir,
+                                                                                epoch=epoch,
+                                                                                num_items=10,)),
+                 CheckpointCallback(vae=vae,
+                                    path=ckpt_dir,
+                                    epoch_interval=args.ckpt_interval,
+                                    restore_training=False,
+                                    restore_path=None),
+                tf.keras.callbacks.TensorBoard(log_dir=log_dir)]
+
+    vae.fit(x=images,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            callbacks=callbacks,
+            shuffle=True,
+            initial_epoch=0)
+

utils/callbacks.py

@@ -0,0 +1,49 @@
+import os.path
+
+import tensorflow as tf
+
+
+class VisualizeCallback(tf.keras.callbacks.Callback):
+    def __init__(self, epoch_interval=1, func=lambda model, epoch: None):
+        super(VisualizeCallback, self).__init__()
+        self.func = func
+        self.epoch_interval = epoch_interval
+
+    def on_epoch_end(self, epoch, logs=None):
+        if epoch % self.epoch_interval == 0 and epoch > 0:
+            self.func(self.model, epoch)
+
+
+class CheckpointCallback(tf.keras.callbacks.Callback):
+    def __init__(self, vae, path, epoch_interval=1, restore_training=False, restore_path=None):
+        super(CheckpointCallback, self).__init__()
+        self.epoch_interval = epoch_interval
+        self.path = path
+        self.vae = vae
+
+        self.ckpt = tf.train.Checkpoint(vae=vae,
+                                        vae_optimizer=vae.optimizer)
+        self.ckpt_manager = tf.train.CheckpointManager(checkpoint=self.ckpt,
+                                                       directory=self.path,
+                                                       max_to_keep=None)
+        self.restore_training = restore_training
+        self.restore_path = restore_path
+        self._saved = False
+
+    def on_epoch_end(self, epoch, logs=None):
+        if epoch % self.epoch_interval == 0 and epoch > 0:
+            self.ckpt_manager.save(checkpoint_number=epoch)
+
+    def on_train_begin(self, logs=None):
+        if self.restore_training:
+            if self.restore_path is None:
+                self.ckpt.restore(self.ckpt_manager.latest_checkpoint).except_partial()
+                print("Resume training from checkpoint ", self.ckpt_manager.latest_checkpoint, "\n")
+            else:
+                self.ckpt.restore(self.restore_path)
+                print("resume training from checkpoint ", self.restore_path, "\n")
+
+    def on_train_end(self, logs=None):
+        weights_path = os.path.join(self.path, "trained-vae")
+        self.ckpt.save(file_prefix=weights_path)
+

utils/dataloader.py

@@ -0,0 +1,64 @@
+import os
+import random
+import numpy as np
+from PIL import Image
+from re import split, compile
+from tensorflow.keras.utils import Sequence
+
+
+def list_filenames(data_path, img_extension='png', filename_prefix=None):
+    if filename_prefix is None:
+        files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension)]
+    else:
+        files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension) and file.startswith(filename_prefix)]
+
+    files_list.sort(key=lambda l: [int(s) if s.isdigit() else s.lower() for s in split(compile(r'(\d+)'), l)])
+    files_list = [os.path.join(data_path, file) for file in files_list]
+    print('Found {} files in {}'.format(len(files_list), data_path))
+    return files_list
+
+
+class Dataset(Sequence):
+    def __init__(self, file_list, batch_size=32, crop_dim=None, resize_dim=None, shuffle=True):
+        self.files_list = file_list
+        self.batch_size = batch_size
+
+        self.crop_dim = crop_dim
+        self.resize_dim = resize_dim
+        self.shuffle = shuffle
+        self.on_epoch_end()
+
+    def __len__(self):
+        return int(np.ceil(len(self.files_list) / self.batch_size))
+
+    def __getitem__(self, idx):
+        batch_files = self.files_list[idx * self.batch_size : (idx + 1) * self.batch_size]
+        images = [self.load_images(f) for f in batch_files]
+        return np.stack(images)
+
+    def on_epoch_end(self):
+        if self.shuffle:
+            np.random.shuffle(self.files_list)
+
+    @staticmethod
+    def center_crop(image, crop_dim):
+        h, w = image.size
+        crop_h, crop_w = crop_dim
+
+        top = max(0, (w - crop_w) // 2)
+        left = max(0, (h - crop_h) // 2)
+        right = min(h - 0, (h + crop_h) // 2)
+        bottom = min(w - 0, (w + crop_w) // 2)
+
+        return image.crop((left, top, right, bottom))
+
+    def load_images(self, filepath):
+        image = Image.open(filepath).convert('RGB')
+        if self.crop_dim:
+            image = self.center_crop(image, crop_dim=self.crop_dim)
+        if self.resize_dim:
+            image = image.resize(self.resize_dim)
+
+        image = np.array(image).astype(np.float32)
+        image = image / 255.0
+        return image

utils/inference_utils.py

@@ -0,0 +1,52 @@
+import os
+import numpy as np
+import tensorflow as tf
+
+from PIL import Image
+
+from model.vae import VAE
+from model import tidev2
+
+
+def init_vae_model(model_name, latent_dim):
+    if model_name == 'tidev2':
+        vae_model = VAE(tidev2.ConvNeXtEncoderTiny(latent_dim=latent_dim),
+                        tidev2.ConvNeXtDecoderTiny(latent_dim=latent_dim)
+                       )
+        return vae_model
+
+
+def load_weights(vae, weights_path):
+    print("Loading weights from {}".format(weights_path))
+    if "ckpt-" in weights_path:
+        ckpt = tf.train.Checkpoint(vae=vae)
+        ckpt.restore(weights_path).expect_partial()
+        return vae
+    if ".TF" in weights_path:
+        vae.load_weights(weights_path, by_name=True)
+        return vae
+
+
+def get_noise_seeded(noise_shape):
+    np.random.seed(0)
+    random_z = np.random.normal(0, 1, noise_shape)
+    return random_z
+
+
+def decode_noise(trained_vae, noise, return_list=False):
+    print("Generating synthetic images ...")
+    pred = trained_vae.decoder.predict(noise, batch_size=1)
+    # print(type(pred), pred.shape, pred.dtype, pred.min(), pred.max())
+    pred *= 255.0
+    # print(type(pred), pred.shape, pred.dtype, pred.min(), pred.max())
+    if return_list:
+        return [img for img in pred]
+    return pred
+
+
+def save_images(save_folder, images):
+    print(f"Saving  synthetic images into {save_folder}")
+    if isinstance(images, list):
+        for i, image in enumerate(images):
+            image = image.astype(np.uint8)
+            Image.fromarray(image).save(os.path.join(save_folder, f"image-{i}.jpg"))

utils/plots.py

@@ -0,0 +1,25 @@
+import imageio
+import numpy as np
+
+
+def visualize_from_latent_space(latent_dim, input_shape, vae, output_path, epoch="final", num_items=10,):
+
+    image_size, _, img_channels = input_shape
+    figure = np.zeros((image_size * num_items, image_size * num_items, 3))
+
+    scale = 1.0
+    grid_x = np.linspace(-scale, scale, num_items)
+    grid_y = np.linspace(-scale, scale, num_items)[::-1]
+
+    np.random.seed(42)
+    for i, yi in enumerate(grid_y):
+        for j, xi in enumerate(grid_x):
+            random_z = np.random.normal(0, 1, (1, latent_dim))
+            x_decoded = vae.decoder.predict(random_z)
+            image = x_decoded[0].reshape(input_shape)
+            figure[i * image_size: (i + 1) * image_size, j * image_size: (j + 1) * image_size, ] = image
+    print(f'Saving collage in {output_path}/decoding-noise-ep{epoch}.jpg')
+    imageio.imsave(f'{output_path}/decoding-noise-ep{epoch}.jpg', (figure * 255).astype('uint8'))
+
+
+