DMAR2dCT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Redução artefatos metálicos (senograma) - RED-Net_.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "-xhFF2iH-WxQ"
      },
      "source": [
        " \n",
        "# Import Tensorflow 2.0\n",
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "import numpy as np\n",
        "import functools\n",
        "from PIL import Image\n",
        "import cv2\n",
        "from numpy import asarray\n",
        "from sklearn.model_selection import train_test_split\n",
        "from tqdm import tqdm\n",
        "from tensorflow.python.framework import ops\n",
        "import timeit\n",
        "from pathlib import Path"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GB7UD-gaeJ7V"
      },
      "source": [
        "#Comando para linkar colab e drive!!\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5e_Wv_XsWjRk"
      },
      "source": [
        "#Celula pra ver GPU\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eII0Ou2gAxGc"
      },
      "source": [
        "import pathlib\n",
        " \n",
        "\n",
        "class dataset_Helper():\n",
        " \n",
        "    def __init__(self, rootDir, trainPerc, testPerc, valPerc):\n",
        "        '''Create dataset helper for each dose reduction factor'''\n",
        "        \n",
        "        if trainPerc + testPerc + valPerc != 1.0:\n",
        "            raise ValueError('Percentages should sum to 1.')\n",
        "           \n",
        "        rootDir2Read = pathlib.Path(rootDir)\n",
        "        \n",
        "        imgFiles = list(rootDir2Read.glob('*SinoMetal*.tiff'))\n",
        "        \n",
        "        \n",
        " \n",
        "        trainAmount =  np.round(len(imgFiles) * trainPerc).astype('int') \n",
        "        testAmount = np.round(len(imgFiles) * testPerc).astype('int') \n",
        "        \n",
        "        imgMetal = []\n",
        "        \n",
        "        for imgFile in imgFiles:\n",
        "            \n",
        "            imgFile = str(imgFile)\n",
        "            \n",
        "            \n",
        " \n",
        "            if 'SinoMetal' in imgFile:\n",
        "            # if imgFile[-14:-5] == 'SinoMetal':\n",
        " \n",
        "                imgMetal.append(imgFile)\n",
        "                \n",
        "        self.noiseData_train = imgMetal[0:trainAmount]\n",
        "        self.noiseData_test  = imgMetal[trainAmount:trainAmount+testAmount]\n",
        "        self.noiseData_val   = imgMetal[trainAmount+testAmount:]\n",
        "                               \n",
        "    def shuffle_dataset(self):\n",
        "        '''Shuffle the TF Dataset on each epoch'''\n",
        "        \n",
        "        perm = list(range(len(self.noiseData_train)))\n",
        "        random.shuffle(perm)\n",
        "        shuffle_array = [self.noiseData_train[index] for index in perm]\n",
        "        self.noiseData_train = shuffle_array\n",
        "        \n",
        "    def create_train_dataset_generator(self, imgBatchSize):\n",
        "        \n",
        "        noiseDataset = self.noiseData_train\n",
        "        \n",
        "        return self.__create_dataset_generator(noiseDataset, imgBatchSize)\n",
        "    \n",
        "    def create_test_dataset_generator(self, imgBatchSize):\n",
        "        \n",
        "        noiseDataset = self.noiseData_test\n",
        "        \n",
        "        return self.__create_dataset_generator(noiseDataset, imgBatchSize)\n",
        "    \n",
        "    def create_val_dataset_generator(self, imgBatchSize):\n",
        "        \n",
        "        noiseDataset = self.noiseData_val\n",
        "        \n",
        "        return self.__create_dataset_generator(noiseDataset, imgBatchSize)\n",
        "        \n",
        "    def __create_dataset_generator(self, noiseDataset, imgBatchSize):\n",
        "        ''' Create a generator to run over the dataset (train, test or val). \n",
        "            The generator get bathches of dicom images'''\n",
        "           \n",
        "        for x in range(0,len(noiseDataset),imgBatchSize):\n",
        "                \n",
        "            batch_noiseDcmFile = noiseDataset[x:x+imgBatchSize]\n",
        "            normal_ROIs_dataset = []\n",
        "            noise_ROIs_dataset = []\n",
        " \n",
        "            for noiseDcmFile in batch_noiseDcmFile:\n",
        "                split = noiseDcmFile.split('/')\n",
        "                split[-1] = split[-1].split('_')[0] + \"_SinoOriginal.tiff\"\n",
        "                normalDcmFile = \"/\".join(split)\n",
        " \n",
        "                # Get noise and normal ROIs from Dicom data\n",
        "                noise_ROIs  = plt.imread(noiseDcmFile).astype('float32') / 255\n",
        "                normal_ROIs = plt.imread(normalDcmFile).astype('float32') / 255\n",
        " \n",
        "                shape = noise_ROIs.shape\n",
        " \n",
        "                noise_ROIs = noise_ROIs.reshape(1, *shape, 1)\n",
        "                normal_ROIs = normal_ROIs.reshape(1, *shape, 1)\n",
        "                # Append data to the ROI dataset list\n",
        "                normal_ROIs_dataset.append(normal_ROIs)\n",
        "                noise_ROIs_dataset.append(noise_ROIs)\n",
        "         \n",
        "            # Stack all the ROIs in one np array\n",
        "            normal_ROIs_dataset =  np.concatenate(normal_ROIs_dataset, axis=0)\n",
        "            noise_ROIs_dataset =  np.concatenate(noise_ROIs_dataset, axis=0)\n",
        "            yield (noise_ROIs_dataset, normal_ROIs_dataset)\n",
        "    \n",
        " \n",
        "dataset = dataset_Helper('/content/drive/My Drive', trainPerc=0.8, testPerc=0.1, valPerc=0.1) \n",
        " \n",
        "roiSize = 180\n",
        "\n",
        "imgBatchSize_train  = 16     # Amount of images to read from train dicom folder\n",
        "imgBatchSize_val = 128     # Amount of images to read from val dicom folder\n",
        "imgBatchSize_test = 128    # Amount of images to read from val dicom folder\n",
        "batch_size = 5\n",
        " \n",
        "# Create a generator to run over the dicom dataset\n",
        "dataset_train_gen = tf.data.Dataset.from_generator(dataset.create_train_dataset_generator,\n",
        "                                              args=(imgBatchSize_train, ), output_types=(tf.float32,tf.float32),\n",
        "                                              output_shapes=((None, roiSize, roiSize, 1), (None, roiSize, roiSize, 1))).prefetch(10)\n",
        " \n",
        "dataset_val_gen = tf.data.Dataset.from_generator(dataset.create_val_dataset_generator,\n",
        "                                              args=(imgBatchSize_val, ), output_types=(tf.float32,tf.float32),\n",
        "                                              output_shapes=((None, roiSize, roiSize, 1), (None, roiSize, roiSize, 1))).prefetch(10)\n",
        " \n",
        "dataset_test_gen = tf.data.Dataset.from_generator(dataset.create_test_dataset_generator,\n",
        "                                              args=(imgBatchSize_test, ), output_types=(tf.float32,tf.float32),\n",
        "                                              output_shapes=((None, roiSize, roiSize, 1), (None, roiSize, roiSize, 1))).prefetch(10)\n",
        " "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u1_xB61X_WEM"
      },
      "source": [
        "class Load_Image():\n",
        "    \n",
        "  def download(self):\n",
        "\n",
        "    self.x_train, self.x_test, self.y_train, self.y_test = train_test_split(array_metal, array_Or, test_size=0, random_state=42)\n",
        "\n",
        "\n",
        "    self.train_dataset_shape = self.x_train.shape\n",
        "    self.test_dataset_shape = self.x_test.shape\n",
        "\n",
        "\n",
        "  def print_information(self):\n",
        "\n",
        "    # Print informations about the DATA dataset.\n",
        "    print(\"There is %d training samples, containing images with shape of: %dx%d and %d channel\" \n",
        "          % (self.train_dataset_shape[0],self.train_dataset_shape[1],self.train_dataset_shape[2],self.train_dataset_shape[3]))\n",
        "    print(\"Train variable shape:\", end='')\n",
        "    print(self.train_dataset_shape)\n",
        "\n",
        "  def pre_process(self):\n",
        "\n",
        "    # Reshape to have 1 channel (last dimension) and normalize\n",
        "    self.x_train = np.expand_dims(self.x_train, axis=-1).astype(np.float32)\n",
        "    self.x_test = np.expand_dims(self.x_test, axis=-1).astype(np.float32)\n",
        "    self.y_train = np.expand_dims(self.y_train, axis=-1).astype(np.float32)\n",
        "    self.y_test = np.expand_dims(self.y_test, axis=-1).astype(np.float32)\n",
        "\n",
        "    # Normalize data\n",
        "    ## https://stats.stackexchange.com/questions/185853/why-do-we-need-to-normalize-the-images-before-we-put-them-into-cnn\n",
        "    self.x_train /= 255\n",
        "    self.x_test /= 255\n",
        "    self.y_train /= 255\n",
        "    self.y_test /= 255\n",
        "\n",
        "\n",
        "    self.train_dataset_shape = self.x_train.shape\n",
        "    self.test_dataset = self.x_test.shape\n",
        "\n",
        "  def train_train(self):\n",
        "\n",
        "    self.x_train_noise = self.x_train \n",
        "    self.x_test_noise = self.x_test \n",
        " \n",
        "  \n",
        "  def create_dataset_iterable(self):\n",
        "\n",
        "    self.train_dataset = tf.data.Dataset.from_tensor_slices((self.x_train_noise, self.y_train))\n",
        "    self.test_dataset = tf.data.Dataset.from_tensor_slices((self.x_test_noise, self.y_test))\n",
        "      \n",
        "  def shuffle_dataset(self, dataset_size):\n",
        "    \n",
        "    return self.train_dataset.shuffle(dataset_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "78ZQ598BH5gb"
      },
      "source": [
        "#Novo!!\n",
        "class autoencoder_model(tf.keras.Model):\n",
        "\n",
        "  def __init__(self):\n",
        "      \n",
        "      super(autoencoder_model, self).__init__()    \n",
        "\n",
        "      self.n_filters = 32\n",
        "      self.img_input_shape = (180, 180, 1) # nRows X nCols X nChannels\n",
        "      \n",
        "      self.encoder = self.build_encoder()\n",
        "      self.decoder = self.build_decoder()\n",
        "      self.autoencoder = self.build_autoencoder()\n",
        "\n",
        "\n",
        "  def build_encoder(self):\n",
        "      \n",
        "      Conv2D = functools.partial(tf.keras.layers.Conv2D, padding='valid', activation=None, strides=1, use_bias=True)      \n",
        "      \n",
        "      imgBatch = tf.keras.layers.Input(shape=self.img_input_shape)\n",
        "      \n",
        "      Conv2D_1 = Conv2D(filters=1*self.n_filters, kernel_size=5)(imgBatch)\n",
        "      ReLU_1 = tf.keras.layers.ReLU()(Conv2D_1)\n",
        "       \n",
        "      Conv2D_2 = Conv2D(filters=1*self.n_filters, kernel_size=5)(ReLU_1)\n",
        "      ReLU_2 = tf.keras.layers.ReLU()(Conv2D_2)\n",
        "       \n",
        "      Conv2D_3 = Conv2D(filters=1*self.n_filters, kernel_size=5)(ReLU_2)\n",
        "      ReLU_3 = tf.keras.layers.ReLU()(Conv2D_3)\n",
        "       \n",
        "      Conv2D_4 = Conv2D(filters=1*self.n_filters, kernel_size=5)(ReLU_3)\n",
        "      ReLU_4 = tf.keras.layers.ReLU()(Conv2D_4)\n",
        "       \n",
        "      Conv2D_5 = Conv2D(filters=1*self.n_filters, kernel_size=5)(ReLU_4)\n",
        "      ReLU_5 = tf.keras.layers.ReLU()(Conv2D_5)\n",
        "      \n",
        "      model = tf.keras.Model(inputs=imgBatch, outputs=[imgBatch, ReLU_2, ReLU_4, ReLU_5], name=\"encoder\")\n",
        "\n",
        "      return model\n",
        "\n",
        "  def build_decoder(self):\n",
        "      '''\n",
        "      https://medium.com/apache-mxnet/transposed-convolutions-explained-with-ms-excel-52d13030c7e8\n",
        "      https://stackoverflow.com/questions/53654310/what-is-the-difference-between-upsampling2d-and-conv2dtranspose-functions-in-ker\n",
        "      '''\n",
        "      \n",
        "      Conv2DTranspose = functools.partial(tf.keras.layers.Conv2DTranspose, padding='valid', activation=None, strides=1, use_bias=True)\n",
        "      ResAdd = tf.keras.layers.Add()\n",
        "      ResWeightedAdd = WeightedAdd(omega=0.5)\n",
        "      \n",
        "      Residual1 = tf.keras.layers.Input(shape=self.img_input_shape)\n",
        "      Residual2 = tf.keras.layers.Input(shape=self.encoder.layers[4].output.shape[1::])   # ReLU_2\n",
        "      Residual3 = tf.keras.layers.Input(shape=self.encoder.layers[8].output.shape[1::])   # ReLU_4\n",
        "\n",
        "      LatentIn  = tf.keras.layers.Input(shape=self.encoder.layers[-1].output.shape[1::])\n",
        "         \n",
        "      # Upscaling convolutions (inverse of encoder)\n",
        "      Conv2DT_1 = Conv2DTranspose(filters=1*self.n_filters, kernel_size=5) (LatentIn)\n",
        "      Add_1 = tf.keras.layers.Add()([Conv2DT_1, Residual3])\n",
        "      ReLU_1 =  tf.keras.layers.ReLU()(Add_1)\n",
        "       \n",
        "      Conv2DT_2 = Conv2DTranspose(filters=1*self.n_filters, kernel_size=5)(ReLU_1)\n",
        "      ReLU_2 = tf.keras.layers.ReLU()(Conv2DT_2)\n",
        "       \n",
        "      Conv2DT_3 = Conv2DTranspose(filters=1*self.n_filters, kernel_size=5)(ReLU_2)\n",
        "      Add_2 = ResAdd([Conv2DT_3, Residual2])\n",
        "      ReLU_3 = tf.keras.layers.ReLU()(Add_2)\n",
        "       \n",
        "      Conv2DT_4 = Conv2DTranspose(filters=1*self.n_filters, kernel_size=5)(ReLU_3)\n",
        "      ReLU_4 = tf.keras.layers.ReLU()(Conv2DT_4)\n",
        "       \n",
        "      Conv2DT_5 = Conv2DTranspose(filters=1, kernel_size=5)(ReLU_4)\n",
        "      Add_3 = ResWeightedAdd([Conv2DT_5, Residual1])\n",
        "      ReLU_5 = tf.keras.layers.ReLU()(Add_3)\n",
        "\n",
        "      model = tf.keras.Model(inputs=[Residual1, Residual2, Residual3, LatentIn], outputs=ReLU_5, name=\"decoder\")\n",
        "       \n",
        "      return model\n",
        "  \n",
        "  def build_autoencoder(self):\n",
        "        \n",
        "      imgBatch = tf.keras.layers.Input(shape=self.img_input_shape)\n",
        "      \n",
        "      Residual1, Residual2, Residual3, LatentIn = self.encoder(imgBatch)\n",
        "      \n",
        "      model = tf.keras.Model(imgBatch, self.decoder(inputs=[Residual1, Residual2, Residual3, LatentIn]), name=\"autoencoder\")\n",
        "      \n",
        "      return model\n",
        "\n",
        "  def call(self, x):\n",
        "      \n",
        "      latent_space = self.encoder(x)\n",
        "      img_recon = self.autoencoder(x)\n",
        "      \n",
        "      return img_recon"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HOE-YtTUVT-H"
      },
      "source": [
        "class WeightedAdd(tf.keras.layers.Layer):\n",
        "    '''\n",
        "    https://datascience.stackexchange.com/questions/56171/how-to-merge-two-cnn-deep-learning-model-using-weighted-sum-and-weighted-product\n",
        "    ''' \n",
        "    \n",
        "    def __init__(self, omega, **kwargs):\n",
        "        \n",
        "        super(WeightedAdd, self).__init__(**kwargs)\n",
        "        \n",
        "        self.omega = 0.5\n",
        "        \n",
        "    def call(self, model_outputs):\n",
        "        \n",
        "        return self.omega * model_outputs[1] + (1 - self.omega) * model_outputs[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YnclH9Atna7r"
      },
      "source": [
        "class perceptualModel(tf.keras.models.Model):\n",
        "  def __init__(self, perceptual_layers):\n",
        "    super(perceptualModel, self).__init__()\n",
        "    self.perceptual_layers = perceptual_layers\n",
        "    self.vgg = self.vgg_layers(perceptual_layers)\n",
        "    self.vgg.trainable = False\n",
        "\n",
        "  def vgg_layers(self, layer_names):\n",
        "    \"\"\" Creates a vgg model that returns a list of intermediate output values.\"\"\"\n",
        "    \n",
        "    # Load our model. Load pretrained VGG, trained on imagenet data\n",
        "    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')\n",
        "    vgg.trainable = False\n",
        "    \n",
        "    outputs = [vgg.get_layer(name).output for name in layer_names]\n",
        "\n",
        "    model = tf.keras.Model([vgg.input], outputs)\n",
        "\n",
        "    return model\n",
        "\n",
        "  def call(self, inputs):\n",
        "    \"Expects float input in [0,1]\"\n",
        "    inputs = inputs*255.0\n",
        "    preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)\n",
        "    perceptual_outputs = self.vgg(preprocessed_input)\n",
        "\n",
        "    perceptual_dict = {perceptual_name:value\n",
        "                  for perceptual_name, value\n",
        "                  in zip(self.perceptual_layers, perceptual_outputs)}\n",
        "    \n",
        "    return perceptual_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eT5X4glonb4T"
      },
      "source": [
        "perceptual_layers = ['block3_conv1', \n",
        "                     'block3_conv2', \n",
        "                     'block3_conv3',\n",
        "                     'block3_conv4']\n",
        "                     \n",
        "perceptualNet = perceptualModel(perceptual_layers)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KrfmvUgonfIS"
      },
      "source": [
        "def loss_function(img_recon, img_without_noise):\n",
        "\n",
        "    # MSE loss\n",
        "    loss_MSE = tf.math.reduce_mean(tf.square(img_recon - img_without_noise),axis=(1,2))\n",
        "\n",
        "    # Perceptual loss\n",
        "    vgg_recon_out = perceptualNet(tf.tile(img_recon,(1,1,1,3)))\n",
        "    vgg_without_noise_out = perceptualNet(tf.tile(img_without_noise,(1,1,1,3)))\n",
        "\n",
        "    perceptual_loss = tf.add_n([tf.reduce_mean((vgg_recon_out[name]-vgg_without_noise_out[name])**2) for name in vgg_recon_out.keys()])\n",
        "    perceptual_loss /= len(perceptual_layers)\n",
        "\n",
        "    # Combination of both losses\n",
        "    loss = loss_MSE + 1.17e-05 * perceptual_loss\n",
        "\n",
        "    return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bnFQD17M_eh8"
      },
      "source": [
        "def create_optimizer(learning_rate):\n",
        "\n",
        "  return tf.keras.optimizers.Adam(learning_rate) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kV2RPL8g_gPB"
      },
      "source": [
        "@tf.function\n",
        "def train_step(autoencoder, img_with_noise , img_without_noise, optimizer):\n",
        "\n",
        "  with tf.GradientTape() as tape:\n",
        "        \n",
        "    # Feed images into autoencoder\n",
        "    img_recon = autoencoder(img_with_noise)\n",
        "        \n",
        "    # Calc the loss\n",
        "    loss = loss_function(img_recon, img_without_noise)\n",
        "\n",
        "    ### Backpropagation ###\n",
        "    # Get the gradients\n",
        "    grads = tape.gradient(loss, autoencoder.trainable_variables)\n",
        "\n",
        "    # Update the weights\n",
        "    optimizer.apply_gradients(zip(grads, autoencoder.trainable_variables))\n",
        "\n",
        "  return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3LyaRI1pzKY2"
      },
      "source": [
        "import os\n",
        "MODEL_SAVE_PATH = \".\"\n",
        "EPS = 1e-4\n",
        "\n",
        "def train_model(autoencoder, train_dataset, validation_dataset, num_epochs, batch_size, learning_rate, patience=25, tol=EPS, save_folder=MODEL_SAVE_PATH):\n",
        "\n",
        "  optimizer = create_optimizer(learning_rate)\n",
        "\n",
        "  train_loss_history = []\n",
        "  validation_loss_history = []\n",
        "\n",
        "  metrics_names = ['train_loss','val_loss'] \n",
        "  \n",
        "  patience_counter = 0\n",
        "  min_loss = None\n",
        "\n",
        "  # Loop on each epoch\n",
        "  for epoch in range(num_epochs):\n",
        "\n",
        "    print(\"\\nepoch {}/{} \".format(epoch+1,num_epochs), end=\"\")\n",
        "\n",
        "    epoch_train_loss = []\n",
        "    for idX, (batch_x, batch_y) in enumerate(train_dataset):\n",
        "      # Train the model\n",
        "      train_loss = train_step(autoencoder, batch_x, batch_y, optimizer)\n",
        "\n",
        "      values=[('train_loss',train_loss)]\n",
        "\n",
        "      # progBar.update(idX*batch_size, values=values) \n",
        "\n",
        "      train_loss_history.append(tf.math.reduce_mean(train_loss))\n",
        "\n",
        "      epoch_train_loss.append(tf.math.reduce_mean(train_loss).numpy())\n",
        "\n",
        "    # Loop on each batch of test dataset for validation\n",
        "    epoch_val_loss = []\n",
        "    for batch_x, batch_y in validation_dataset:\n",
        "\n",
        "      # Foward image through the network\n",
        "      img_recon = autoencoder(batch_x)\n",
        "\n",
        "      # Calc the loss\n",
        "      val_loss = loss_function(img_recon, batch_y)\n",
        "\n",
        "      validation_loss_history.append(tf.math.reduce_mean(val_loss))\n",
        "\n",
        "      epoch_val_loss.append(tf.math.reduce_mean(val_loss).numpy())\n",
        "\n",
        "    val_loss_reduced_mean = np.mean(epoch_val_loss)\n",
        "    if min_loss is None or abs(val_loss_reduced_mean - min_loss) > tol:\n",
        "      min_loss = val_loss_reduced_mean\n",
        "      autoencoder.save_weights(os.path.join(save_folder, 'model_checkpoint'))\n",
        "      patience_counter = 0\n",
        "    else:\n",
        "      patience_counter +=1 \n",
        "\n",
        "\n",
        "    print(\" Train loss:\", epoch_train_loss[-1], \"| Val. loss:\", epoch_val_loss[-1])\n",
        "\n",
        "    if patience_counter > patience:\n",
        "      print(f\"STOPING AT EPOCH {epoch}: PATIENCE EXCEEDED.\")\n",
        "      num_epochs = epoch + 1\n",
        "      break\n",
        "  \n",
        "  train_loss = [train_loss_history[x].numpy() for x in range(0,len(train_loss_history),len(train_loss_history)//num_epochs)]\n",
        "  val_loss = [validation_loss_history[x].numpy() for x in range(0,len(validation_loss_history),len(validation_loss_history)//num_epochs)]\n",
        "\n",
        "\n",
        "  return train_loss, val_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W8duPe05_olX"
      },
      "source": [
        "autoencoder = autoencoder_model()\n",
        "autoencoder.encoder.summary()\n",
        "autoencoder.decoder.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IJ4gTK-B_rMR"
      },
      "source": [
        "# Training hyperparameters\n",
        "num_epochs = 500\n",
        "batch_size = 16\n",
        "learning_rate = 2e-4\n",
        "patience = 20  \n",
        "tol = 1e-3\n",
        "#save_folder = \"./checkpoint\"\n",
        "# Train the model \n",
        "train_loss, val_loss = train_model(\n",
        "    autoencoder, \n",
        "    dataset_train_gen,\n",
        "    dataset_val_gen,\n",
        "    num_epochs,\n",
        "    batch_size,\n",
        "    learning_rate,\n",
        "    patience=patience,\n",
        "    tol=tol,\n",
        "    save_folder=save_folder\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8rItXEcA49oU"
      },
      "source": [
        "# Loading a saved model - Rodar quando tiver model_checkpoint salvo!\n",
        "autoencoder.load_weights(os.path(save_folder, \"model_checkpoint\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tOxbN2LU_szi"
      },
      "source": [
        "# Plot loss results\n",
        "\n",
        "plt.figure()\n",
        "print(train_loss)\n",
        "print(val_loss)\n",
        "plt.plot(train_loss, label=\"train_loss\")\n",
        "plt.plot(val_loss, label=\"val_loss\")\n",
        "plt.title(\"Results\")\n",
        "plt.xlabel(\"Epoch #\")\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.legend(loc=\"uppdataset_test_gener right\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C6XBzkbesWj7"
      },
      "source": [
        "import os\n",
        "from os.path import join\n",
        "\n",
        "N = 1\n",
        "\n",
        "image_save_folder = \"./image_save\"\n",
        "\n",
        "if not os.path.exists(image_save_folder):\n",
        "  os.mkdir(image_save_folder)\n",
        "\n",
        "for batch_x, batch_y in dataset_test_gen:\n",
        "\n",
        "  for i, (img_metal, img_or) in enumerate(zip(batch_x, batch_y)):\n",
        "    \n",
        "\n",
        "    print(tf.reshape(img_metal, (1, *img_metal.shape)))\n",
        "    img_recon = autoencoder(tf.reshape(img_metal, (1, *img_metal.shape)))\n",
        "\n",
        "    print(img_recon.numpy().squeeze())\n",
        "    plt.figure(figsize=(5, 5))\n",
        "    ax_metal = plt.gca()\n",
        "    ax_metal.imshow(img_metal.numpy().squeeze(),'gray')\n",
        "    ax_metal.set_title(\"Metal\", fontsize=15)\n",
        "    \n",
        "    plt.figure(figsize=(5, 5))\n",
        "    ax_orig = plt.gca()\n",
        "    ax_orig.imshow(img_or.numpy().squeeze(),'gray')\n",
        "    ax_orig.set_title(\"Original\", fontsize=15)\n",
        "    \n",
        "    plt.figure(figsize=(5, 5))\n",
        "    ax_rec = plt.gca()\n",
        "    ax_rec.imshow(img_recon.numpy().squeeze(),'gray')\n",
        "    ax_rec.set_title(\"Reconstructed\", fontsize=15)\n",
        "  \n",
        "\n",
        "\n",
        "    plt.imsave(join(image_save_folder, f\"{i}_SinoCorrigido.png\"), img_recon.numpy().squeeze(), cmap='gray')\n",
        "\n",
        "\n",
        "    if i > N:\n",
        "      break\n",
        "    \n",
        "    plt.savefig(join(image_save_folder, \"\"), cmap='gray')\n",
        "  break\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}