IA_DNN_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "IA_DNN_GAN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MedicalImageAnalysisTutorials/DeepLearning4All/blob/main/IA_DNN_GAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UlkWTXM3WCcd"
      },
      "source": [
        "# **Generative Adversarial Network (GAN)**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42gVqgVhXATV"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MhUEkuFSDOcQ"
      },
      "source": [
        "A generative adversarial network (GAN) is a class of machine learning frameworks designed by Ian Goodfellow and his colleagues in 2014.\n",
        "Generative Adversarial Networks, or GANs for short, are an approach to generative modeling using deep learning methods, such as convolutional neural networks.\n",
        "\n",
        "Generative modeling is an unsupervised learning task in machine learning that involves automatically discovering and learning the regularities or patterns in \n",
        "input data in such a way that the model can be used to generate or output new examples that plausibly could have been drawn from the original dataset.\n",
        "\n",
        "\n",
        "GAN has two parts:\n",
        "\n",
        "The generator learns to generate plausible data. The generated instances become negative training examples for the discriminator.\n",
        "The discriminator learns to distinguish the generator's fake data from real data. The discriminator penalizes the generator for producing implausible results.\n",
        "\n",
        "1.When training begins, the generator produces obviously fake data, and the discriminator quickly learns to tell that it's fake.\n",
        "\n",
        "2.As training progresses, the generator gets closer to producing output that can fool the discriminator.\n",
        "\n",
        "3.Finally, if generator training goes well, the discriminator gets worse at telling the difference between real and fake. It starts to classify fake data as real, and its accuracy decreases."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CNwoZlHaypNQ"
      },
      "source": [
        "### **Applications**\n",
        "\n",
        "1.**Data Augmentation**: Data augmentation in data analysis are techniques used to increase the amount of data by adding slightly modified copies of already existing data or newly created synthetic data from existing data.\n",
        "Medical Application\n",
        "\n",
        "![](https://machinelearningmastery.com/wp-content/uploads/2019/04/Example-of-Vector-Arithmetic-for-GAN-Generated-Faces.png)\n",
        "\n",
        "2.**Image Translation**\n",
        "\n",
        "\n",
        "![](https://machinelearningmastery.com/wp-content/uploads/2019/06/Example-of-Sketches-to-Color-Photographs-with-pix2pix.png)\n",
        "\n",
        "\n",
        "**Some more applications are**:\n",
        "\n",
        "\n",
        "Medical Application\n",
        "\n",
        "Generate Examples for Image Datasets\n",
        "\n",
        "Generate Photographs of Human Faces\n",
        "\n",
        "Generate Realistic Photographs\n",
        "\n",
        "Generate Cartoon Characters\n",
        "\n",
        "Image-to-Image Translation\n",
        "\n",
        "Text-to-Image Translation\n",
        "\n",
        "Semantic-Image-to-Photo Translation\n",
        "\n",
        "Face Frontal View Generation\n",
        "\n",
        "Generate New Human Poses\n",
        "\n",
        "Photos to Emojis\n",
        "\n",
        "Photograph Editing\n",
        "\n",
        "Face Aging\n",
        "\n",
        "Photo Blending\n",
        "\n",
        "Super Resolution\n",
        "\n",
        "Photo Inpainting\n",
        "\n",
        "Clothing Translation\n",
        "\n",
        "Video Prediction\n",
        "\n",
        "3D Object Generation\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T4qztqDcFDwo"
      },
      "source": [
        "###**Architecture**\n",
        "A picture of the whole system:\n",
        "\n",
        "![](https://developers.google.com/machine-learning/gan/images/gan_diagram.svg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-8MwoxteFJjq"
      },
      "source": [
        "###**Loss function**\n",
        "\n",
        "**Discriminator:**\n",
        "$$\n",
        "D_{loss} = \\frac{1}{m} \\sum_{i=1}^{m} [\\log{D(x^{(i)})} + \\log{(1-D(G(z^{(i)})))}]\n",
        "$$\n",
        "\n",
        "**Generator:**\n",
        "$$\n",
        "G_{loss} = \\frac{1}{m} \\sum_{i=1}^{m} \\log{(1-D(G(z^{(i)})))}\n",
        "$$\n",
        "\n",
        "where $D$ is the discriminator and $G$ is the generator. $x^{(i)}$ is the $i$th ground truth data and $z^{(i)}$ is the $i$th generated data(random input). $m$ is the total number of data.\n",
        "\n",
        "The discriminator is to find the mistake and the generator wants to fake the discriminator.\n",
        "\n",
        "If $D()$ is 1 means the discriminator thinks the data is real. if $D()$ is 0 means it is fake. But note that $\\log{0}$  is negative infinity.\n",
        "\n",
        "We want to max the loss function for discriminator and minimize the loss function for generator. Loss function close to 0 is max. Loss function close to negative infinity is min.\n",
        "\n",
        "Combining the two losses together,\n",
        "$$\\min_{G}\\max_{D} V(D,G) = E_x[\\log{D(x)}] + E_z[\\log{(1-D(G(z)))}]$$\n",
        "\n",
        "where $E_x$ is the expected value over all real data instances. $E_z$ is the expected value over all random inputs to the generator.\n",
        "\n",
        "In practice, we change the generator loss function to \n",
        "$$ G_{loss} = \\max_{G} E_z[\\log{D(G(z))}]$$\n",
        "\n",
        "Discriminator loss function:\n",
        "$$\\max_{D} V(D,G) = E_x[\\log{D(x)}] + E_z[\\log{(1-D(G(z)))}]$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0beVL6O7hco"
      },
      "source": [
        "# Notebook Setup "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NooQ5tm85ZDb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8fee70ae-5d05-4181-d3a1-5e8a4a5d64ea"
      },
      "source": [
        "# Preparing 2d datset from medical images (ct to mri and mri to ct) \n",
        "# download medical image dataset use only t2 and ct\n",
        "# unzip, load to slicer then save as nrrd\n",
        "# reset the scene and load the nrrd files\n",
        "# register (low to high) e.g. t2  to ct\n",
        "# save with same name of the moving image and add \"_reg.nrrd\"\n",
        "# this is our 3d dataset\n",
        "# during training: \n",
        "#   img = read using simpleitk \n",
        "#   imgArr = convert to array \n",
        "#   this array is now your 2d datset e.g. 28, 512,512 \n",
        "# Setup \n",
        "usePt = 1\n",
        "doInstall =1\n",
        "if doInstall:\n",
        "  !pip install SimpleITK\n",
        "# !pip install albumentations\n",
        "!pip install albumentations==0.4.6\n",
        "import albumentations \n",
        "from albumentations.pytorch import ToTensorV2\n",
        "import os, random, time, math\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import albumentations as A\n",
        "import cv2 \n",
        "import SimpleITK as sitk \n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from PIL import Image\n",
        "\n",
        "if usePt == 0:\n",
        "  from tensorflow.keras import layers\n",
        "  from tensorflow.keras import datasets, layers, models\n",
        "\n",
        "if usePt == 1:\n",
        "  import torch\n",
        "  import torch.nn as nn\n",
        "  import torch.optim as optim\n",
        "  import torchvision\n",
        "  # import torchvision.datasets as datasets\n",
        "  from torchvision import datasets\n",
        "  from torch.utils.data import DataLoader, Dataset\n",
        "  import torchvision.transforms as transforms\n",
        "  from torch.utils.tensorboard import SummaryWriter  # to print to tensorboard\n",
        "  from torchvision import transforms\n",
        "\n",
        "\n",
        "# to reproduce the same results given same input\n",
        "np.random.seed(1)               \n"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: SimpleITK in /usr/local/lib/python3.7/dist-packages (2.1.1)\n",
            "Requirement already satisfied: albumentations==0.4.6 in /usr/local/lib/python3.7/dist-packages (0.4.6)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from albumentations==0.4.6) (3.13)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from albumentations==0.4.6) (1.4.1)\n",
            "Requirement already satisfied: imgaug>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from albumentations==0.4.6) (0.4.0)\n",
            "Requirement already satisfied: opencv-python>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from albumentations==0.4.6) (4.1.2.30)\n",
            "Requirement already satisfied: numpy>=1.11.1 in /usr/local/lib/python3.7/dist-packages (from albumentations==0.4.6) (1.19.5)\n",
            "Requirement already satisfied: Shapely in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (1.7.1)\n",
            "Requirement already satisfied: imageio in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (2.4.1)\n",
            "Requirement already satisfied: scikit-image>=0.14.2 in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (0.16.2)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (3.2.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (1.15.0)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.7/dist-packages (from imgaug>=0.4.0->albumentations==0.4.6) (7.1.2)\n",
            "Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image>=0.14.2->imgaug>=0.4.0->albumentations==0.4.6) (2.6.3)\n",
            "Requirement already satisfied: PyWavelets>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image>=0.14.2->imgaug>=0.4.0->albumentations==0.4.6) (1.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->imgaug>=0.4.0->albumentations==0.4.6) (2.8.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->imgaug>=0.4.0->albumentations==0.4.6) (2.4.7)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->imgaug>=0.4.0->albumentations==0.4.6) (0.10.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->imgaug>=0.4.0->albumentations==0.4.6) (1.3.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01cHwK5kmDxi"
      },
      "source": [
        "# GAN Hyper Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hiz_0QlsRDg5"
      },
      "source": [
        "# notebook parameters\n",
        "if usePt == 1:\n",
        "  device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "doAug = 0\n",
        "GANs = ['simple', 'DCGAN', 'WGAN', 'P2P'] # 0, 1, 2, 3\n",
        "GANID = 3\n",
        "datasetID = 5  # 1:minst is selected by default, for cifar10 use 2\n",
        "               # 5: mr-ct dataset\n",
        "NNID      = 3  # 1:NN is by default, for DNN use 2,or 3, for 3D use 4  \n",
        "number_of_classes = 10  # each datasets have 10 classes\n",
        "showSamples = 1\n",
        "UseOurDatasets = 0\n",
        "new_size = (64,64)\n",
        "# GAN parameters\n",
        "LEARNING_RATE = 2e-4  # could also use two lrs, one for gen and one for dsc\n",
        "BATCH_SIZE = 32\n",
        "IMAGE_SIZE = 64\n",
        "CHANNELS_IMG = 1 # channels\n",
        "Z_DIM = 100 # \n",
        "NUM_EPOCHS = 5 \n",
        "FEATURES_DISC = 64 \n",
        "FEATURES_GEN = 64 \n",
        "CRITIC_ITERATIONS = 5\n",
        "WEIGHT_CLIP = 0.01\n",
        "\n",
        "# if you have large GPU memory you can combine the images to batches \n",
        "# for faster training.\n",
        "# It is good to try different values\n",
        "batch_size = 32 # you can try larger batch size e.g. 1024 * 6"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cB5haVaRy2sg"
      },
      "source": [
        "# Dataset preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zs8ZXeOxudH2"
      },
      "source": [
        "### Datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2tqEEtuQufSr"
      },
      "source": [
        "# $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ Using MNIST\n",
        "from albumentations.pytorch import ToTensorV2\n",
        "def datasetReturn(GANID):\n",
        "  import torchvision.transforms as transforms\n",
        "  if GANID == 0:\n",
        "    transforms = transforms.Compose(\n",
        "        [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)),]\n",
        "    )\n",
        "  else:\n",
        "    transforms = transforms.Compose([transforms.Resize(IMAGE_SIZE),transforms.ToTensor(),transforms.Normalize([0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]),])\n",
        "\n",
        "  dataset = datasets.MNIST(root=\"dataset/\", transform=transforms, download=True)\n",
        "  return dataset\n",
        "  \n",
        "both_transform = A.Compose(\n",
        "    [A.Resize(width=256, height=256),], additional_targets={\"image0\": \"image\"},\n",
        ")\n",
        "\n",
        "transform_only_input = A.Compose(\n",
        "    [\n",
        "        A.HorizontalFlip(p=0.5),\n",
        "        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0,),\n",
        "        ToTensorV2(),\n",
        "    ]\n",
        ")\n",
        "transform_only_mask = A.Compose(\n",
        "    [\n",
        "        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0,),\n",
        "        ToTensorV2(),\n",
        "    ]\n",
        ")\n",
        "\n",
        "\n",
        "# $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ Using our own datasets\n",
        "class MedicalDataset(torch.utils.data.Dataset):\n",
        "    def __init__(self, root_dir, transform=None): \n",
        "        self.root_dir = root_dir   \n",
        "        self.transform = transform \n",
        "        self.images = os.listdir(self.root_dir)\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.images)\n",
        "    \n",
        "    def __getitem__(self,index):#\n",
        "        image_index = self.images[index]# Using index to get image\n",
        "        img_path = os.path.join(self.root_dir, image_index)# get path\n",
        "        img = io.imread(img_path)# read image\n",
        "        label = img_path.split('\\\\')[-1].split('.')[0]# get label\n",
        "        sample = {'image':img,'label':label}\n",
        "\n",
        "\n",
        "class MapDataset(Dataset):\n",
        "    def __init__(self, root_dir):\n",
        "        self.root_dir = root_dir\n",
        "        self.list_files = os.listdir(self.root_dir)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.list_files)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        img_file = self.list_files[index]\n",
        "        img_path = os.path.join(self.root_dir, img_file)\n",
        "        image = np.array(Image.open(img_path))\n",
        "        input_image = image[:, :600, :]\n",
        "        target_image = image[:, 600:, :]\n",
        "\n",
        "        augmentations = both_transform(image=input_image, image0=target_image)\n",
        "        input_image = augmentations[\"image\"]\n",
        "        target_image = augmentations[\"image0\"]\n",
        "\n",
        "        input_image = transform_only_input(image=input_image)[\"image\"]\n",
        "        target_image = transform_only_mask(image=target_image)[\"image\"]\n",
        "\n",
        "        return input_image, target_image\n",
        "\n",
        "\n",
        "# if UseOurDa\n",
        "# data = AnimalData('E:/Python Project/PyTorch/dogs-vs-cats/train',transform=None)\n",
        "\n",
        "\n",
        "# # if __name__=='__main__':\n",
        "#     data = AnimalData('E:/Python Project/PyTorch/dogs-vs-cats/train',transform=None)\n",
        "#     dataloader = DataLoader(data,batch_size=BATCH_SIZE,shuffle=True)\n",
        "#     for i_batch,batch_data in enumerate(dataloader):\n",
        "#         print(i_batch)\n",
        "#         print(batch_data['image'].size())\n",
        "#         print(batch_data['label'])\n",
        "\n"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s8b7mb_FjhwH"
      },
      "source": [
        "# !wget --no-check-certificate https://www.kaggle.com/vikramtiwari/pix2pix-dataset?select=cityscapes -O ctcscapes-datasets.tar\n",
        "# !unzip ct-mr-datasets.zip"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 252
        },
        "id": "Dzj-Tvek4V8s",
        "outputId": "2ffe4774-4859-4fbb-bf37-0289558abb4e"
      },
      "source": [
        "def read3DMI(imgPath,doNormalisation=1, new_size=[]):\n",
        "    img = sitk.ReadImage(imgPath)\n",
        "    # print(img.GetSize())\n",
        "    imga = sitk.GetArrayFromImage(img) # sitk reverse the 3rd dimension\n",
        "    # print(img.GetSize())\n",
        "    # print(imga.shape)\n",
        "    if len(new_size)>0:      \n",
        "        imgaNew= [] #np.zeros(new_size)\n",
        "        # resize image\n",
        "        for i in range (imga.shape[0]):\n",
        "            s = imga[i,:,:]\n",
        "            v = cv2.resize(s, (64,64), interpolation = cv2.INTER_AREA)\n",
        "            imgaNew.append(v) #cv2.resize(imga[i,:,:], new_size, interpolation = cv2.INTER_AREA))\n",
        "        imga = np.array(imgaNew)  \n",
        "\n",
        "    if doNormalisation:\n",
        "       imgMax = np.max(imga) \n",
        "       imgMin = np.min(imga) \n",
        "       imga = (imga - imgMin) / (imgMax-imgMin)\n",
        "            \n",
        "    return img, imga\n",
        "\n",
        "# minst dataset\n",
        "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
        "class_names = range(10)\n",
        "if datasetID==2:\n",
        "    # cifar10 dataset\n",
        "    # The CIFAR10 dataset contains 60,000 color images in 10 classes, \n",
        "    # with 6,000 images in each class.\n",
        "    # The dataset is divided into 50,000 training images and 10,000 testing images.\n",
        "    # The classes are mutually exclusive and there is no overlap between them.\n",
        "\n",
        "    (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()\n",
        "    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',\n",
        "               'dog', 'frog', 'horse', 'ship', 'truck']\n",
        "    if NNID==4:\n",
        "        #TODO: fix this \n",
        "        showSamples =0\n",
        "        \n",
        "        x_train = x_train.reshape(-1,32*32*3) # (32 * 32 * 3)        \n",
        "        x_train = np.resize(x_train,(x_train.shape[0],15,15,15))        \n",
        "        x_train = x_train.reshape(-1,15,15,15)\n",
        "        x_test = x_test.reshape(-1,32*32*3) # (32 * 32 * 3)\n",
        "        x_test = np.resize(x_test,(x_test.shape[0],15,15,15))\n",
        "        x_test = x_test.reshape(-1,15,15,15)\n",
        "        x_train  =  x_train[..., np.newaxis] # np.reshape(x_train, (-1, h,w,1))\n",
        "        y_train  =  y_train[..., np.newaxis] # np.reshape(y_train, (-1, h,w,1))\n",
        "        x_test   =  x_test[..., np.newaxis]  # np.reshape(x_test,  (-1, h,w,1))\n",
        "        y_test   =  y_test[..., np.newaxis]  # np.reshape(y_test,  (-1, h,w,1))\n",
        "\n",
        "\n",
        "        print(x_train.shape)\n",
        "        print(x_test.shape)\n",
        "elif datasetID==5:\n",
        "    #  !wget --no-check-certificate https://cloud.uni-koblenz.de/s/GS8YKpC83A6jLr8/download/mri_ct_images.zip -O ct-mr-datasets.zip\n",
        "    #  !unzip ct-mr-datasets.zip\n",
        "     fnms = sorted(os.listdir(\"Registration_image\"))\n",
        "     ctFnms = [x for x in fnms if \"ct\" in x]\n",
        "     mrFnms = [x for x in fnms if \"reg\" in x]\n",
        "     print(\"fnms : \",len(fnms))\n",
        "     print(\"ctFnms　: \",len(mrFnms))\n",
        "     print(ctFnms)\n",
        "     print(mrFnms)\n",
        "     ct_dataset = []\n",
        "     mr_dataset = []\n",
        "     num_ct = 0\n",
        "     num_mri = 0\n",
        "    #  img,imgsCT =  read3DMI(\"/content/Registration_image/\"+ctFnms[0],doNormalisation=1, new_size=(64,64))\n",
        "    #  patient_000_ct \n",
        "    #  patient_000_mr_T2_reg\n",
        "     for i in range(len(mrFnms)):\n",
        "        img,imgsCT =  read3DMI(\"/content/Registration_image/\"+ctFnms[i],doNormalisation=1, new_size=new_size)\n",
        "        img,imgsMR =  read3DMI(\"/content/Registration_image/\"+mrFnms[i],doNormalisation=1, new_size=new_size)\n",
        "        # 29, 512 x 512   ---> num_2d, 64,64        \n",
        "        # 40, 256 x 256   ---> num_2d, 64,64\n",
        "        # ...\n",
        "        # \n",
        "        # sum(num_2d),64,64     :X=ct,Y=mri\n",
        "        #\n",
        "        num_ct  += imgsCT.shape[0] \n",
        "        num_mri += imgsCT.shape[0] \n",
        "        ct_dataset.extend(imgsCT)\n",
        "        mr_dataset.extend(imgsMR)\n",
        "     print(num_mri) # 238\n",
        "     print(num_ct)\n",
        "     ct_dataset = np.array(ct_dataset)\n",
        "     mr_dataset = np.array(mr_dataset)\n",
        "     ct_dataset = ct_dataset[..., np.newaxis]\n",
        "     mr_dataset = mr_dataset[..., np.newaxis]\n",
        "\n",
        "     train_CT  = np.array(ct_dataset)\n",
        "     train_mri = np.array(mr_dataset)\n",
        "      # Convert to tensor \n",
        "     train_CT  = torch.from_numpy(train_CT)\n",
        "     train_mri = torch.from_numpy(train_mri)\n",
        "     print(train_CT.shape)\n",
        "      # get size \n",
        "     h, w = 64, 64\n",
        "    #  print(ok)\n",
        "    # check for rgb \n",
        "    #dataset = torch.from_numpy()\n",
        "\n",
        "\n",
        "\n",
        "try:\n",
        "    # number of channels\n",
        "    c =  x_train.shape[3]\n",
        "except:\n",
        "    # number of channels\n",
        "    c =  1\n",
        "    # if there is no number of channels, add 1\n",
        "    x_train  =  x_train[..., np.newaxis] # np.reshape(x_train, (-1, h,w,1))\n",
        "    y_train  =  y_train[..., np.newaxis] # np.reshape(y_train, (-1, h,w,1))\n",
        "    x_test   =  x_test[..., np.newaxis]  # np.reshape(x_test,  (-1, h,w,1))\n",
        "    y_test   =  y_test[..., np.newaxis]  # np.reshape(y_test,  (-1, h,w,1))\n",
        "\n",
        "\n",
        "# Reserve 10,000 samples for validation.\n",
        "x_val = x_train[-10000:]\n",
        "y_val = y_train[-10000:]\n",
        "x_train = x_train[:-10000]\n",
        "y_train = y_train[:-10000]\n",
        "h, w = 64, 64\n",
        "number_of_pixels = h * w * c\n",
        "\n",
        "\n",
        "print(\"dataset shape   : \",x_train.shape)\n",
        "print(\"number of images: \",x_train.shape[0])\n",
        "print(\"image size      : \",x_train[0].shape)\n",
        "print(\"image data type : \",type(x_train[0][0][0][0]))\n",
        "print(\"image max  value: \",np.max(x_train[0]))\n",
        "print(\"image min  value: \",np.min(x_train[0]))\n",
        "if c==1:\n",
        "   print(\"gray or binary image (not color image)\")\n",
        "elif c==3:\n",
        "   print(\"rgb color image (or probably non-color image represented with 3 channels)\")\n",
        "\n",
        "\n",
        "# display sample images \n",
        "if showSamples:\n",
        "    plt.figure(figsize=(10,10))\n",
        "    for i in range(25):\n",
        "        plt.subplot(5,5,i+1)\n",
        "        plt.xticks([])\n",
        "        plt.yticks([])\n",
        "        plt.grid(False)\n",
        "        #plt.imshow(x_train[i])\n",
        "        print('kkkkkkkkkkkkkkkkkkkkkkkkkk', type(x_train[i]))\n",
        "        plt.imshow(cv2.cvtColor(x_train[i], cv2.COLOR_BGR2RGB))\n",
        "\n",
        "        # The CIFAR labels happen to be arrays, \n",
        "        # which is why you need the extra index\n",
        "        if datasetID==1:\n",
        "            plt.xlabel(y_train[i])\n",
        "        elif datasetID==2:\n",
        "            plt.xlabel(class_names[y_train[i][0]])\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "# normalisation\n",
        "x_train = np.array([ x/255.0 for x in x_train])\n",
        "x_val   = np.array([ x/255.0 for x in x_val])\n",
        "x_test  = np.array([ x/255.0 for x in x_test])\n",
        "#y_train = y_train.astype(np.float32)\n",
        "\n",
        "# for NN we need 1D \n",
        "if NNID ==1:\n",
        "   x_train = np.reshape(x_train, (-1, number_of_pixels))\n",
        "   x_val   = np.reshape(x_val,  (-1, number_of_pixels))\n",
        "   x_test  = np.reshape(x_test , (-1, number_of_pixels))\n",
        "print(x_train.shape,y_train.shape)\n",
        "# x_train = np.reshape(x_train, (50000, 64, 64, 1))\n",
        "# x_train = np.resize(x_train[:,])\n",
        "\n",
        "# Prepare the training dataset.\n",
        "print(x_train.shape,y_train.shape)\n",
        "# print(ok)\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)\n",
        "#train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)\n",
        "\n",
        "# Prepare the validation dataset.\n",
        "val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))\n",
        "val_dataset = val_dataset.batch(batch_size)\n",
        "\n",
        "# Prepare the test dataset.\n",
        "tst_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "tst_dataset = tst_dataset.batch(batch_size)"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "error",
          "ename": "FileNotFoundError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-15-951d297d688a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     55\u001b[0m     \u001b[0;31m#  !wget --no-check-certificate https://cloud.uni-koblenz.de/s/GS8YKpC83A6jLr8/download/mri_ct_images.zip -O ct-mr-datasets.zip\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0;31m#  !unzip ct-mr-datasets.zip\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 57\u001b[0;31m      \u001b[0mfnms\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msorted\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlistdir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Registration_image\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     58\u001b[0m      \u001b[0mctFnms\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mx\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfnms\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;34m\"ct\"\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m      \u001b[0mmrFnms\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mx\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfnms\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;34m\"reg\"\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'Registration_image'"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kshgRoz2sqZG"
      },
      "source": [
        "fnms = sorted(os.listdir(\"Registration_image\"))\n",
        "ctFnms = [x for x in fnms if 'ct' in x]\n",
        "ctFnms"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHTqF_g8soX1"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bP0E_Pzs04y_"
      },
      "source": [
        "## Dataset augmentation\n",
        "\n",
        "It is important to train the model on different variations of the dataset. It is also important to have large datset for training.\n",
        "\n",
        "Using dataset augmentation helps to achieve both of the above goals. From one image, one can generate hundred thousands of images using image transformation.\n",
        "\n",
        "The image transformation could be [spatial transform]() or point transform where we move the points of the image to new locations e.g. shifting, flipping, and/or rotating the imag. \n",
        "\n",
        "Another type of transformation is intensity transform or pixel transform where we change the color values of the pixels in the image e.g. invert the color, add more brightness or darkness. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jb2KdoTR07L2"
      },
      "source": [
        "def imagePixelTransforms(img):    \n",
        "    images = []\n",
        "    # let's make 3 simple transformations\n",
        "    img1   = 1.0- img # invert color\n",
        "    img2   = img +0.3 # more brightness\n",
        "    img3   = img -0.3 # more darkness\n",
        "    images = np.array([img1,img2,img3])\n",
        "    images = [ img.reshape(img.shape) for img in images]\n",
        "\n",
        "    # plt.figure() ;    plt.imshow(img)\n",
        "    # plt.figure() ;    plt.imshow(img1)\n",
        "    # plt.figure() ;    plt.imshow(img2)\n",
        "    # plt.figure() ;    plt.imshow(img3)    \n",
        "    return images\n",
        "\n",
        "def imagePointTransforms(img):\n",
        "    images = []\n",
        "    # let's make 3 simple transformations\n",
        "    # Perform the rotation\n",
        "    center  = (img.shape[0] / 2, img.shape[1] / 2)\n",
        "    sz      = (img.shape[1], img.shape[0])\n",
        "    tMatrix = cv2.getRotationMatrix2D(center, 45, 1)\n",
        "    img1 = cv2.warpAffine(img, tMatrix, sz)\n",
        "    img1 = img1[...,np.newaxis] if img1.shape !=img.shape else img1\n",
        "    tMatrix = cv2.getRotationMatrix2D(center, 90, 1)\n",
        "    img2 = cv2.warpAffine(img, tMatrix, sz)\n",
        "    img2 = img2[...,np.newaxis] if img2.shape !=img.shape else img2\n",
        "    tMatrix = cv2.getRotationMatrix2D(center, 270, 1)\n",
        "    img3 = cv2.warpAffine(img, tMatrix, sz)\n",
        "    img3 = img3[...,np.newaxis] if img3.shape !=img.shape else img3\n",
        "\n",
        "    images = np.array([img1,img2,img3])\n",
        "    # plt.figure() ;    plt.imshow(img)\n",
        "    # plt.figure() ;    plt.imshow(img1)\n",
        "    # plt.figure() ;    plt.imshow(img2)\n",
        "    # plt.figure() ;    plt.imshow(img3)    \n",
        "    #plt.figure() ;    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
        "    # plt.figure() ;    plt.imshow(cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))\n",
        "    return images\n",
        "\n",
        "\n",
        "# define a function for sitk transform\n",
        "def resample(img_array, transform):\n",
        "    # Output image Origin, Spacing, Size, Direction are taken from the reference\n",
        "    # image in this call to Resample\n",
        "    image = sitk.GetImageFromArray(img_array)\n",
        "    reference_image = image\n",
        "    interpolator = sitk.sitkCosineWindowedSinc\n",
        "    default_value = 100.0\n",
        "    resampled_img = sitk.Resample(image, reference_image, transform,\n",
        "                         interpolator, default_value)\n",
        "    resampled_array = sitk.GetArrayFromImage(resampled_img)\n",
        "    return resampled_array\n",
        "\n",
        "def affine_rotate(transform, degrees):\n",
        "    parameters = np.array(transform.GetParameters())\n",
        "    new_transform = sitk.AffineTransform(transform)\n",
        "    dimension =3 \n",
        "    matrix = np.array(transform.GetMatrix()).reshape((dimension,dimension))\n",
        "    radians = -np.pi * degrees / 180.\n",
        "    rotation = np.array([[1  ,0,0], \n",
        "                         [0, np.cos(radians), -np.sin(radians)],\n",
        "                         [0, np.sin(radians), np.cos(radians)]]\n",
        "                        )\n",
        "    new_matrix = np.dot(rotation, matrix)\n",
        "    new_transform.SetMatrix(new_matrix.ravel())\n",
        "    return new_transform\n",
        "\n",
        "\n",
        "def imagePoint3DTransforms(img):\n",
        "    #print(\"imagePoint3DTransforms\")\n",
        "    images = []\n",
        "    # let's make 3 simple transformations\n",
        "    # Perform the rotation\n",
        "    # In SimpleITK resampling convention, the transformation maps points \n",
        "    # from the fixed image to the moving image,\n",
        "    # so inverse of the transform is applied\n",
        "\n",
        "    center = (img.shape[0] /2, img.shape[1] /2,img.shape[1] /2)\n",
        "    rotation_around_center = sitk.AffineTransform(3)\n",
        "    rotation_around_center.SetCenter(center)\n",
        "    \n",
        "    rotation_around_center = affine_rotate(rotation_around_center, -45)\n",
        "    img1 = resample(img, rotation_around_center)\n",
        "\n",
        "    rotation_around_center = affine_rotate(rotation_around_center, -90)\n",
        "    img2 = resample(img, rotation_around_center)\n",
        "\n",
        "    rotation_around_center = affine_rotate(rotation_around_center, -90)\n",
        "    img3 = resample(img, rotation_around_center)\n",
        "\n",
        "    images = np.array([img1,img2,img3])\n",
        "    # plt.figure() ;    plt.imshow(img)\n",
        "    # plt.figure() ;    plt.imshow(img1)\n",
        "    # plt.figure() ;    plt.imshow(img2)\n",
        "    # plt.figure() ;    plt.imshow(img3)    \n",
        "    #plt.figure() ;    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
        "    # plt.figure() ;    plt.imshow(cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))\n",
        "    return images\n",
        "\n",
        "def doAugmentation(images,labels,batch_size):\n",
        "    # input is an image or a batch e.g. list of images \n",
        "    # get numpy arrays from the tensor    \n",
        "    images = images.numpy()\n",
        "    labels = labels.numpy()\n",
        "    # if 1d convert back to 2d\n",
        "    #print(images.shape)\n",
        "    rgb = 0 ; is3d = 0\n",
        "    if len(images.shape) == 2:\n",
        "       try: \n",
        "          img2d_shape = int(math.sqrt(images.shape[1])) # gray or binary image\n",
        "          images =images.reshape(-1,img2d_shape,img2d_shape)\n",
        "       except:\n",
        "          try: \n",
        "            img2d_shape = int(math.sqrt(images.shape[1]/3)) # rgb image\n",
        "            images =images.reshape(-1,img2d_shape,img2d_shape,3)\n",
        "            rgb = 1  \n",
        "          except:\n",
        "            pass  \n",
        "            # img3d_shape = int(math.sqrt(images.shape[1]/3)) # rgb image\n",
        "            # images =images.reshape(-1,img2d_shape,img2d_shape,3)\n",
        "            # is3d = 1  \n",
        "\n",
        "\n",
        "\n",
        "    x_outputs = [] ; y_outputs = []\n",
        "    i = 0\n",
        "    for img in images:\n",
        "        #print(\"-------------------------\", i ,\"--------------------\")\n",
        "        if NNID==4:\n",
        "           img = img.squeeze() \n",
        "        # from each images we generate 6 images\n",
        "        # 64 batch will generate 448\n",
        "        x_outputs.extend([img])\n",
        "        imgs1 = imagePoint3DTransforms(img)\n",
        "        imgs2 = imagePixelTransforms(img)\n",
        "        #if not rgb:\n",
        "           #imgs1 = np.array( x[...,np.newaxis] for x in imgs1 if len(x.shape)<3) \n",
        "           #imgs2 = np.array( x[...,np.newaxis] for x in imgs2 if len(x.shape)<3)\n",
        "        x_outputs.extend(imgs1) # 3 images\n",
        "        x_outputs.extend(imgs2) # 3 images\n",
        "        # print(img.shape)\n",
        "        # print(imgs1[0].shape)\n",
        "        # print(imgs2[0].shape)\n",
        "        # assign the same label to all transformed images\n",
        "        for j in range ( len(imgs1) +len(imgs2)+1):\n",
        "            y_outputs.extend([labels[i]])\n",
        "\n",
        "        i = i +1\n",
        "    x_outputs = np.array(x_outputs)\n",
        "    if NNID==4:\n",
        "       x_outputs = np.array([x[...,np.newaxis] for x in x_outputs])\n",
        "    y_outputs = np.array(y_outputs)\n",
        "\n",
        "    if (not rgb) and (NNID==1):\n",
        "       x_outputs = np.reshape(x_outputs, (-1,img2d_shape*img2d_shape,1))\n",
        "    elif (rgb) and (NNID==1):\n",
        "       x_outputs = np.reshape(x_outputs, (-1,img2d_shape*img2d_shape*3))   \n",
        "\n",
        "    new_train_dataset = tf.data.Dataset.from_tensor_slices((x_outputs, y_outputs))\n",
        "    new_train_dataset = new_train_dataset.shuffle(buffer_size=1024).batch(batch_size)\n",
        "\n",
        "    return new_train_dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PYHsww4UfDKX"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lZ75bXEepHft"
      },
      "source": [
        "##  1 Simple GAN model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pcr20SeXQc37"
      },
      "source": [
        "this is basic one-hiden layer neural network without concolution involved"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hTWNUEHbYAmp"
      },
      "source": [
        "### Implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R7fxsV4_9EDK"
      },
      "source": [
        "if usePt == 0:\n",
        "  discriminator_tf_simple = keras.Sequential(\n",
        "      [\n",
        "      keras.Input(shape = (784,)), # 784  #28\n",
        "      layers.Dense(128),\n",
        "      layers.LeakyReLU(0.01),\n",
        "      layers.Dense(1, activation='sigmoid'),\n",
        "      #  layers.Sigmoid(),\n",
        "      ],\n",
        "      name = \"discriminator_tf\",\n",
        "  )\n",
        "\n",
        "  generator_tf_simple = keras.Sequential(\n",
        "      [\n",
        "          keras.Input(shape=(64,)),\n",
        "          layers.Dense(256),\n",
        "          # layers.Reshape((8, 8, 128)),\n",
        "          layers.LeakyReLU(alpha=0.01),\n",
        "          layers.Dense(784, activation='tanh'), # 28x28\n",
        "          # layers.Conv2D(3, kernel_size=5, padding=\"same\", activation=\"sigmoid\"),\n",
        "      ],\n",
        "      name=\"generator_tf\",\n",
        "  )\n",
        "\n",
        "  criterion = keras.losses.BinaryCrossentropy()\n",
        "  opt_gen = keras.optimizers.Adam(3e-4)\n",
        "  opt_dis = keras.optimizers.Adam(3e-4)\n",
        "else:\n",
        "    # Pytorch\n",
        "    class SimpleDiscriminator(nn.Module):\n",
        "        def __init__(self, in_features):\n",
        "            super().__init__()\n",
        "            print(in_features)\n",
        "            self.disc = nn.Sequential(\n",
        "                nn.Linear(in_features, 128),\n",
        "                nn.LeakyReLU(0.01),\n",
        "                nn.Linear(128, 1),\n",
        "                nn.Sigmoid(),\n",
        "            )\n",
        "\n",
        "        def forward(self, x):\n",
        "            return self.disc(x)\n",
        "\n",
        "    class SimpleGenerator(nn.Module):\n",
        "        def __init__(self, z_dim, img_dim):\n",
        "            super().__init__()\n",
        "            self.gen = nn.Sequential(\n",
        "                nn.Linear(z_dim, 256),\n",
        "                nn.LeakyReLU(0.01),\n",
        "                nn.Linear(256, img_dim),\n",
        "                nn.Tanh(),  # normalize inputs to [-1, 1] so make outputs [-1, 1]\n",
        "            )\n",
        "\n",
        "        def forward(self, x):\n",
        "            return self.gen(x)\n",
        "\n",
        "    # optimisers             \n",
        "    # opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE)\n",
        "    # opt_gen  = optim.Adam(gen.parameters(), lr=LEARNING_RATE)\n",
        "    \n",
        "    # loss function\n",
        "    criterion = nn.BCELoss()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMC0RuBvPNoT"
      },
      "source": [
        "## 2 DCGAN(Deep Convolutional Generative Adversarial Networks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "71VKyXPdXyw-"
      },
      "source": [
        "### Introduction\n",
        "\n",
        "In recent years, supervised learning with convolutional networks (CNNs) has seen huge adoption in computer vision applications. Comparatively, unsupervised learning with CNNs has received less attention. In this work we hope to help bridge the gap between the success of CNNs for supervised learning and unsupervised learning. We introduce a class of CNNs called deep convolutional generative adversarial networks (DCGANs), that have certain architectural constraints, and demonstrate that they are a strong candidate for unsupervised learning.\n",
        "\n",
        "DCGAN uses convolutional and convolutional-transpose layers in the generator and discriminator, respectively. Here the discriminator consists of strided convolution layers, batch normalization layers, and LeakyRelu as activation function. It takes a 3x64x64 input image. The generator consists of convolutional-transpose layers, batch normalization layers, and ReLU activations. The output will be a 3x64x64 RGB image\n",
        "\n",
        "![](https://editor.analyticsvidhya.com/uploads/2665314.png)\n",
        "\n",
        "Figure:DCGAN generator used for LSUN scene modeling. A 100 dimensional uniform distribution Z is projected to a small spatial extent convolutional representation with many feature maps.\n",
        "A series of four fractionally-strided convolutions (in some recent papers, these are wrongly called\n",
        "deconvolutions) then convert this high level representation into a 64 × 64 pixel image. Notably, no\n",
        "fully connected or pooling layers are used.\n",
        "\n",
        "Referance:https://arxiv.org/pdf/1511.06434.pdf\n",
        "\n",
        "paper link: https://arxiv.org/abs/1511.06434"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W2qqukK8X2OQ"
      },
      "source": [
        "### Implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FbE2usPQYHsk"
      },
      "source": [
        "## tensorflow"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dERth2SAYH-V"
      },
      "source": [
        "#----------------------------------------------------------\n",
        "#\n",
        "#DCGAN model definition\n",
        "#\n",
        "#----------------------------------------------------------\n",
        "## pytorch\n",
        "\"\"\"\n",
        "Discriminator and Generator implementation from DCGAN paper\n",
        "\"\"\"\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "if usePt == 0:\n",
        "  discriminator_tf_dc = keras.Sequential(\n",
        "      [\n",
        "        keras.Input(shape = (64,64,1)), # 784  #28\n",
        "        layers.Conv2D(FEATURES_DISC, kernel_size = 4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(0.2),\n",
        "        layers.Conv2D(FEATURES_DISC*2, kernel_size = 4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(0.2),\n",
        "        layers.Conv2D(FEATURES_DISC*4, kernel_size = 4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(0.2),\n",
        "        layers.Conv2D(FEATURES_DISC*8, kernel_size = 4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(0.2),layers.Flatten(),\n",
        "        layers.Dropout(0.2),\n",
        "        layers.Dense(1, activation=\"sigmoid\"),\n",
        "      ],\n",
        "      name = \"discriminator_tf_dc\",\n",
        "  )\n",
        "  latent_dim = 128\n",
        "  generator_tf_dc = keras.Sequential(\n",
        "    [\n",
        "        keras.Input(shape=(latent_dim,)),\n",
        "        layers.Dense(8 * 8 * 128),\n",
        "        layers.Reshape((8, 8, 128)),\n",
        "        layers.Conv2DTranspose(128, kernel_size=4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(alpha=0.2),\n",
        "        layers.Conv2DTranspose(256, kernel_size=4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(alpha=0.2),\n",
        "        layers.Conv2DTranspose(512, kernel_size=4, strides=2, padding=\"same\"),\n",
        "        layers.LeakyReLU(alpha=0.2),\n",
        "        layers.Conv2D(1, kernel_size=5, padding=\"same\", activation=\"sigmoid\"),\n",
        "    ],\n",
        "    name=\"generator\",\n",
        ")\n",
        "else:\n",
        "  class DCGAN_Discriminator(nn.Module):\n",
        "      def __init__(self, channels_img, features_d):\n",
        "          super(DCGAN_Discriminator, self).__init__()\n",
        "          self.disc = nn.Sequential(\n",
        "              # input: N x channels_img x 64 x 64\n",
        "              nn.Conv2d(\n",
        "                  channels_img, features_d, kernel_size=4, stride=2, padding=1\n",
        "              ),\n",
        "              nn.LeakyReLU(0.2),\n",
        "              # _block(in_channels, out_channels, kernel_size, stride, padding)\n",
        "              self._block(features_d, features_d * 2, 4, 2, 1),\n",
        "              self._block(features_d * 2, features_d * 4, 4, 2, 1),\n",
        "              self._block(features_d * 4, features_d * 8, 4, 2, 1),\n",
        "              # After all _block img output is 4x4 (Conv2d below makes into 1x1)\n",
        "              nn.Conv2d(features_d * 8, 1, kernel_size=4, stride=2, padding=0),\n",
        "              nn.Sigmoid(),\n",
        "          )\n",
        "\n",
        "      def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "          return nn.Sequential(\n",
        "              nn.Conv2d(\n",
        "                  in_channels,\n",
        "                  out_channels,\n",
        "                  kernel_size,\n",
        "                  stride,\n",
        "                  padding,\n",
        "                  bias=False,\n",
        "              ),\n",
        "              #nn.BatchNorm2d(out_channels),\n",
        "              nn.LeakyReLU(0.2),\n",
        "          )\n",
        "\n",
        "      def forward(self, x):\n",
        "          return self.disc(x)\n",
        "\n",
        "\n",
        "  class DCGAN_Generator(nn.Module):\n",
        "      def __init__(self, channels_noise, channels_img, features_g):\n",
        "          super(DCGAN_Generator, self).__init__()\n",
        "          self.net = nn.Sequential(\n",
        "              # Input: N x channels_noise x 1 x 1\n",
        "              self._block(channels_noise, features_g * 16, 4, 1, 0),  # img: 4x4\n",
        "              self._block(features_g * 16, features_g * 8, 4, 2, 1),  # img: 8x8\n",
        "              self._block(features_g * 8, features_g * 4, 4, 2, 1),  # img: 16x16\n",
        "              self._block(features_g * 4, features_g * 2, 4, 2, 1),  # img: 32x32\n",
        "              nn.ConvTranspose2d(\n",
        "                  features_g * 2, channels_img, kernel_size=4, stride=2, padding=1\n",
        "              ),\n",
        "              # Output: N x channels_img x 64 x 64\n",
        "              nn.Tanh(),\n",
        "          )\n",
        "\n",
        "      def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "          return nn.Sequential(\n",
        "              nn.ConvTranspose2d(\n",
        "                  in_channels,\n",
        "                  out_channels,\n",
        "                  kernel_size,\n",
        "                  stride,\n",
        "                  padding,\n",
        "                  bias=False,\n",
        "              ),\n",
        "              #nn.BatchNorm2d(out_channels),\n",
        "              nn.ReLU(),\n",
        "          )\n",
        "\n",
        "      def forward(self, x):\n",
        "          return self.net(x)\n",
        "\n",
        "  # def test():\n",
        "  #     N, in_channels, H, W = 8, 3, 64, 64\n",
        "  #     noise_dim = 100\n",
        "  #     x = torch.randn((N, in_channels, H, W))\n",
        "  #     disc = Discriminator(in_channels, 8)\n",
        "  #     assert disc(x).shape == (N, 1, 1, 1), \"Discriminator test failed\"\n",
        "  #     gen = Generator(noise_dim, in_channels, 8)\n",
        "  #     z = torch.randn((N, noise_dim, 1, 1))\n",
        "  #     assert gen(z).shape == (N, in_channels, H, W), \"Generator test failed\"\n",
        "\n",
        "\n",
        "# test()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UAI_u2rSPOLn"
      },
      "source": [
        "## 3  WGAN(Wasserstein GAN)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RmO92iA6X6tJ"
      },
      "source": [
        "### Introduction\n",
        "\n",
        "The Wasserstein Generative Adversarial Network, or W-GAN, is an extension to the generative adversarial network that both improves the stability when training the model and provides a loss function that correlates with the quality of generated images.\n",
        "\n",
        "The development of the WGAN has a dense mathematical motivation, although in practice requires only a few minor modifications to the established standard deep convolutional generative adversarial network, or DCGAN.\n",
        "\n",
        "It is an extension of the GAN that seeks an alternate way of training the generator model to better approximate the distribution of data observed in a given training dataset.\n",
        "\n",
        "Instead of using a discriminator to classify or predict the probability of generated images as being real or fake, the WGAN changes or replaces the discriminator model with a critic that scores the realness or fakeness of a given image.\n",
        "\n",
        "This change is motivated by a theoretical argument that training the generator should seek a minimization of the distance between the distribution of the data observed in the training dataset and the distribution observed in generated examples.\n",
        "\n",
        "The benefit of the WGAN is that the training process is more stable and less sensitive to model architecture and choice of hyperparameter configurations. Perhaps most importantly, the loss of the discriminator appears to relate to the quality of images created by the generator.\n",
        "\n",
        "\n",
        "Referance Paper:\n",
        "https://machinelearningmastery.com/how-to-code-a-wasserstein-generative-adversarial-network-wgan-from-scratch/\n",
        "\n",
        "Resources and papers: \\\\\n",
        "https://www.alexirpan.com/2017/02/22/... \\\\\n",
        "https://arxiv.org/abs/1701.07875 \\\\\n",
        "https://arxiv.org/abs/1704.00028 \\\\"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MXLoxA9GX6tK"
      },
      "source": [
        "### Implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Krfggxg0YP6U"
      },
      "source": [
        "## tensorflow\n",
        "##-------------------------------------------\n",
        "\n",
        "##-------------------------------------------discriminator\n",
        "\n",
        "def discriminator_tf_wgan(input, is_train):\n",
        "    with tf.variable_scope('discriminator') as scope:\n",
        "        # 64*64*64\n",
        "        conv1 = tf.layers.conv2d(input, 64, kernel_size=5, strides=2, padding=\"SAME\",\n",
        "                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                 name='conv1')\n",
        "        act1 = leaky_relu(conv1, n='act1')\n",
        " \n",
        "        # 32*32*128\n",
        "        conv2 = tf.layers.conv2d(act1, 128, kernel_size=5, strides=2, padding=\"SAME\",\n",
        "                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                 name='conv2')\n",
        "        bn2 = tf.contrib.layers.batch_norm(conv2, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn2')\n",
        "        act2 = leaky_relu(bn2, n='act2')\n",
        " \n",
        "        # 16*16*256\n",
        "        conv3 = tf.layers.conv2d(act2, 256, kernel_size=5, strides=2, padding=\"SAME\",\n",
        "                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                 name='conv3')\n",
        "        bn3 = tf.contrib.layers.batch_norm(conv3, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn3')\n",
        "        act3 = leaky_relu(bn3, n='act3')\n",
        " \n",
        "        # 8*8*512\n",
        "        conv4 = tf.layers.conv2d(act3, 512, kernel_size=5, strides=2, padding=\"SAME\",\n",
        "                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                 name='conv4')\n",
        "        bn4 = tf.contrib.layers.batch_norm(conv4, is_training=is_train, epsilon=1e-5, decay=0.9, updates_collections=None,\n",
        "                                           scope='bn4')\n",
        "        act4 = leaky_relu(bn4, n='act4')\n",
        " \n",
        "        # start from act4\n",
        "        dim = int(np.prod(act4.get_shape()[1:]))\n",
        "        fc1 = tf.reshape(act4, shape=[-1, dim], name='fc1')\n",
        "        w2 = tf.get_variable('w2', shape=[fc1.shape[-1], 1], dtype=tf.float32,\n",
        "                             initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
        "        b2 = tf.get_variable('b2', shape=[1], dtype=tf.float32,\n",
        "                             initializer=tf.constant_initializer(0.0))\n",
        "        # wgan不适用sigmoid\n",
        "        logits = tf.add(tf.matmul(fc1, w2), b2, name='logits')\n",
        " \n",
        "        return logits\n",
        "##---------------------------------------------------------generator\n",
        "def generator_tf_wgan(input, random_dim, is_train, reuse=False):\n",
        "    with tf.variable_scope('generator') as scope:\n",
        "        if reuse:\n",
        "            scope.reuse_variables()\n",
        "        w1 = tf.get_variable('w1', shape=[random_dim, 4 * 4 * 512], dtype=tf.float32,\n",
        "                             initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
        "        b1 = tf.get_variable('b1', shape=[512 * 4 * 4], dtype=tf.float32,\n",
        "                             initializer=tf.constant_initializer(0.0))\n",
        "        flat_conv1 = tf.add(tf.matmul(input, w1), b1, name='flat_conv1')\n",
        " \n",
        "        # 4*4*512\n",
        "        conv1 = tf.reshape(flat_conv1, shape=[-1, 4, 4, 512], name='conv1')\n",
        "        bn1 = tf.contrib.layers.batch_norm(conv1, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn1')\n",
        "        act1 = tf.nn.relu(bn1, name='act1')\n",
        " \n",
        "        # 8*8*256\n",
        "        conv2 = tf.layers.conv2d_transpose(act1, 256, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\n",
        "                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                           name='conv2')\n",
        "        bn2 = tf.contrib.layers.batch_norm(conv2, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn2')\n",
        "        act2 = tf.nn.relu(bn2, name='act2')\n",
        " \n",
        "        # 16*16*128\n",
        "        conv3 = tf.layers.conv2d_transpose(act2, 128, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\n",
        "                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                           name='conv3')\n",
        "        bn3 = tf.contrib.layers.batch_norm(conv3, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn3')\n",
        "        act3 = tf.nn.relu(bn3, name='act3')\n",
        " \n",
        "        # 32*32*64\n",
        "        conv4 = tf.layers.conv2d_transpose(act3, 64, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\n",
        "                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                           name='conv4')\n",
        "        bn4 = tf.contrib.layers.batch_norm(conv4, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn4')\n",
        "        act4 = tf.nn.relu(bn4, name='act4')\n",
        " \n",
        "        # 64*64*32\n",
        "        conv5 = tf.layers.conv2d_transpose(act4, 32, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\n",
        "                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                           name='conv5')\n",
        "        bn5 = tf.contrib.layers.batch_norm(conv5, is_training=is_train, epsilon=1e-5, decay=0.9,\n",
        "                                           updates_collections=None, scope='bn5')\n",
        "        act5 = tf.nn.relu(bn5, name='act5')\n",
        " \n",
        "        # 128*128*3\n",
        "        conv6 = tf.layers.conv2d_transpose(act5, image_channel, kernel_size=[5, 5], strides=[2, 2], padding=\"SAME\",\n",
        "                                           kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
        "                                           name='conv6')\n",
        " \n",
        "        act6 = tf.nn.tanh(conv6, name='act6')\n",
        " \n",
        "        return act6"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "96Cdogm-YP6V"
      },
      "source": [
        "#----------------------------------------------------------\n",
        "#\n",
        "#WGAN model definition\n",
        "#\n",
        "#----------------------------------------------------------\n",
        "## pytorch\n",
        "\"\"\"\n",
        "Discriminator and Generator implementation from DCGAN paper,\n",
        "with removed Sigmoid() as output from Discriminator (and therefor\n",
        "it should be called critic)\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "class WGAN_Discriminator(nn.Module):\n",
        "    def __init__(self, channels_img, features_d):\n",
        "        super(WGAN_Discriminator, self).__init__()\n",
        "        \n",
        "        self.disc = nn.Sequential(  \n",
        "            # input: N x channels_img x 64 x 64\n",
        "            nn.Conv2d(channels_img, features_d, kernel_size=4, stride=2, padding=1),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            \n",
        "            # _block(in_channels, out_channels, kernel_size, stride, padding)\n",
        "            self._block(features_d, features_d * 2, 4, 2, 1),\n",
        "            self._block(features_d * 2, features_d * 4, 4, 2, 1),\n",
        "            self._block(features_d * 4, features_d * 8, 4, 2, 1),\n",
        "            # After all _block img output is 4x4 (Conv2d below makes into 1x1)\n",
        "            nn.Conv2d(features_d * 8, 1, kernel_size=4, stride=2, padding=0),\n",
        "        )\n",
        "        # ??? output shape? \n",
        "\n",
        "    def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "        return nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_channels,\n",
        "                out_channels,\n",
        "                kernel_size,\n",
        "                stride,\n",
        "                padding,\n",
        "                bias=False,            ),\n",
        "            #???\n",
        "            nn.InstanceNorm2d(out_channels, affine=True),\n",
        "            nn.LeakyReLU(0.2),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.disc(x)\n",
        "\n",
        "\n",
        "class WGAN_Generator(nn.Module):\n",
        "    def __init__(self, channels_noise, channels_img, features_g):\n",
        "        super(WGAN_Generator, self).__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            # Input: N x channels_noise x 1 x 1\n",
        "            self._block(channels_noise, features_g * 16, 4, 1, 0),  # img: 4x4\n",
        "            self._block(features_g * 16, features_g * 8, 4, 2, 1),  # img: 8x8\n",
        "            self._block(features_g * 8, features_g * 4, 4, 2, 1),  # img: 16x16\n",
        "            self._block(features_g * 4, features_g * 2, 4, 2, 1),  # img: 32x32\n",
        "            nn.ConvTranspose2d(\n",
        "                features_g * 2, channels_img, kernel_size=4, stride=2, padding=1\n",
        "            ),\n",
        "            # Output: N x channels_img x 64 x 64\n",
        "            nn.Tanh(),\n",
        "        )\n",
        "\n",
        "    def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "        return nn.Sequential(\n",
        "            nn.ConvTranspose2d(\n",
        "                in_channels,\n",
        "                out_channels,\n",
        "                kernel_size,\n",
        "                stride,\n",
        "                padding,\n",
        "                bias=False,\n",
        "            ),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jf5CWR56XScB"
      },
      "source": [
        "## 4 Pix2Pix GAN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IzypynDOX7JH"
      },
      "source": [
        "### Introduction\n",
        "\n",
        "Image-to-image translation is the controlled conversion of a given source image to a target image.\n",
        "\n",
        "An example might be the conversion of black and white photographs to color photographs.\n",
        "\n",
        "Image-to-image translation is a challenging problem and often requires specialized models and loss functions for a given translation task or dataset.\n",
        "\n",
        "The Pix2Pix GAN is a general approach for image-to-image translation. It is based on the conditional generative adversarial network, where a target image is generated, conditional on a given input image. In this case, the Pix2Pix GAN changes the loss function so that the generated image is both plausible in the content of the target domain, and is a plausible translation of the input image.\n",
        "\n",
        "Referance:https://machinelearningmastery.com/a-gentle-introduction-to-pix2pix-generative-adversarial-network/\n",
        "\n",
        "Paper link: \\\\\n",
        "https://arxiv.org/abs/1611.07004"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yBn2XjF-X7JH"
      },
      "source": [
        "### Implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LuJTJdSqYRK-"
      },
      "source": [
        "## tensorflow\n",
        "#----------------------------------------------- discriminator\n",
        "class Discriminator_tf_p2p(object):\n",
        "    def __init__(self, inputs, is_training, stddev=0.02, center=True, scale=True, reuse=None):\n",
        "        self._is_training = is_training\n",
        "        self._stddev = stddev\n",
        "\n",
        "        with tf.variable_scope('D', initializer=tf.truncated_normal_initializer(stddev=self._stddev), reuse=reuse):\n",
        "            self._center = center\n",
        "            self._scale = scale\n",
        "            self._prob = 0.5 # constant from pix2pix paper\n",
        "            self._inputs = inputs\n",
        "            self._discriminator = self._build_discriminator(inputs, reuse=reuse)\n",
        "\n",
        "    def _build_layer(self, name, inputs, k, bn=True, use_dropout=False):\n",
        "        layer = dict()\n",
        "        with tf.variable_scope(name):\n",
        "            layer['filters'] = tf.get_variable('filters', [4, 4, get_shape(inputs)[-1], k])\n",
        "            layer['conv'] = tf.nn.conv2d(inputs, layer['filters'], strides=[1, 2, 2, 1], padding='SAME')\n",
        "            layer['bn'] = batch_norm(layer['conv'], center=self._center, scale=self._scale, training=self._is_training) if bn else layer['conv']\n",
        "            layer['dropout'] = tf.nn.dropout(layer['bn'], self._prob) if use_dropout else layer['bn']\n",
        "            layer['fmap'] = lkrelu(layer['dropout'], slope=0.2)\n",
        "        return layer\n",
        "\n",
        "    def _build_discriminator(self, inputs, reuse=None):\n",
        "        discriminator = dict()\n",
        "\n",
        "        # C64-C128-C256-C512 -> PatchGAN\n",
        "        discriminator['l1'] = self._build_layer('l1', inputs, 64, bn=False)\n",
        "        discriminator['l2'] = self._build_layer('l2', discriminator['l1']['fmap'], 128)\n",
        "        discriminator['l3'] = self._build_layer('l3', discriminator['l2']['fmap'], 256)\n",
        "        discriminator['l4'] = self._build_layer('l4', discriminator['l3']['fmap'], 512)\n",
        "        with tf.variable_scope('l5'):\n",
        "            l5 = dict()\n",
        "            l5['filters'] = tf.get_variable('filters', [4, 4, get_shape(discriminator['l4']['fmap'])[-1], 1])\n",
        "            l5['conv'] = tf.nn.conv2d(discriminator['l4']['fmap'], l5['filters'], strides=[1, 1, 1, 1], padding='SAME')\n",
        "            l5['bn'] = batch_norm(l5['conv'], center=self._center, scale=self._scale, training=self._is_training)\n",
        "            l5['fmap'] = tf.nn.sigmoid(l5['bn'])\n",
        "            discriminator['l5'] = l5\n",
        "        return discriminator\n",
        "#----------------------------------------------- generator\n",
        "class Generator_tf_p2p(object):\n",
        "    def __init__(self, inputs, is_training, ochan, stddev=0.02, center=True, scale=True, reuse=None):\n",
        "        self._is_training = is_training\n",
        "        self._stddev = stddev\n",
        "        self._ochan = ochan\n",
        "        with tf.variable_scope('G', initializer=tf.truncated_normal_initializer(stddev=self._stddev), reuse=reuse):\n",
        "            self._center = center\n",
        "            self._scale = scale\n",
        "            self._prob = 0.5 # constant from pix2pix paper\n",
        "            self._inputs = inputs\n",
        "            self._encoder = self._build_encoder(inputs)\n",
        "            self._decoder = self._build_decoder(self._encoder)\n",
        "\n",
        "    def _build_encoder_layer(self, name, inputs, k, bn=True, use_dropout=False):\n",
        "        layer = dict()\n",
        "        with tf.variable_scope(name):\n",
        "            layer['filters'] = tf.get_variable('filters', [4, 4, get_shape(inputs)[-1], k])\n",
        "            layer['conv'] = tf.nn.conv2d(inputs, layer['filters'], strides=[1, 2, 2, 1], padding='SAME')\n",
        "            layer['bn'] = batch_norm(layer['conv'], center=self._center, scale=self._scale, training=self._is_training) if bn else layer['conv']\n",
        "            layer['dropout'] = tf.nn.dropout(layer['bn'], self._prob) if use_dropout else layer['bn']\n",
        "            layer['fmap'] = lkrelu(layer['dropout'], slope=0.2)\n",
        "        return layer\n",
        "\n",
        "    def _build_encoder(self, inputs):\n",
        "        encoder = dict()\n",
        "\n",
        "        # C64-C128-C256-C512-C512-C512-C512-C512\n",
        "        with tf.variable_scope('encoder'):\n",
        "            encoder['l1'] = self._build_encoder_layer('l1', inputs, 64, bn=False)\n",
        "            encoder['l2'] = self._build_encoder_layer('l2', encoder['l1']['fmap'], 128)\n",
        "            encoder['l3'] = self._build_encoder_layer('l3', encoder['l2']['fmap'], 256)\n",
        "            encoder['l4'] = self._build_encoder_layer('l4', encoder['l3']['fmap'], 512)\n",
        "            encoder['l5'] = self._build_encoder_layer('l5', encoder['l4']['fmap'], 512)\n",
        "            encoder['l6'] = self._build_encoder_layer('l6', encoder['l5']['fmap'], 512)\n",
        "            encoder['l7'] = self._build_encoder_layer('l7', encoder['l6']['fmap'], 512)\n",
        "            encoder['l8'] = self._build_encoder_layer('l8', encoder['l7']['fmap'], 512)\n",
        "        return encoder\n",
        "\n",
        "    def _build_decoder_layer(self, name, inputs, output_shape_from,use_dropout=False):\n",
        "        layer = dict()\n",
        "\n",
        "        with tf.variable_scope(name):\n",
        "            output_shape = tf.shape(output_shape_from)\n",
        "            layer['filters'] = tf.get_variable('filters', [4, 4, get_shape(output_shape_from)[-1], get_shape(inputs)[-1]])\n",
        "            layer['conv'] = tf.nn.conv2d_transpose(inputs, layer['filters'], output_shape=output_shape, strides=[1, 2, 2, 1], padding='SAME')\n",
        "            layer['bn'] = batch_norm(tf.reshape(layer['conv'], output_shape), center=self._center, scale=self._scale, training=self._is_training)\n",
        "            layer['dropout'] = tf.nn.dropout(layer['bn'], self._prob) if use_dropout else layer['bn']\n",
        "            layer['fmap'] = tf.nn.relu(layer['dropout'])\n",
        "        return layer\n",
        "\n",
        "    def _build_decoder(self, encoder):\n",
        "        decoder = dict()\n",
        "\n",
        "        # CD512-CD1024-CD1024-C1024-C1024-C512-C256-C128\n",
        "        with tf.variable_scope('decoder'): # U-Net\n",
        "            decoder['dl1'] = self._build_decoder_layer('dl1', encoder['l8']['fmap'], output_shape_from=encoder['l7']['fmap'], use_dropout=True)\n",
        "\n",
        "            # fmap_concat represent skip connections\n",
        "            fmap_concat = tf.concat([decoder['dl1']['fmap'], encoder['l7']['fmap']], axis=3)\n",
        "            decoder['dl2'] = self._build_decoder_layer('dl2', fmap_concat, output_shape_from=encoder['l6']['fmap'], use_dropout=True)\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl2']['fmap'], encoder['l6']['fmap']], axis=3)\n",
        "            decoder['dl3'] = self._build_decoder_layer('dl3', fmap_concat, output_shape_from=encoder['l5']['fmap'], use_dropout=True)\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl3']['fmap'], encoder['l5']['fmap']], axis=3)\n",
        "            decoder['dl4'] = self._build_decoder_layer('dl4', fmap_concat, output_shape_from=encoder['l4']['fmap'])\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl4']['fmap'], encoder['l4']['fmap']], axis=3)\n",
        "            decoder['dl5'] = self._build_decoder_layer('dl5', fmap_concat, output_shape_from=encoder['l3']['fmap'])\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl5']['fmap'], encoder['l3']['fmap']], axis=3)\n",
        "            decoder['dl6'] = self._build_decoder_layer('dl6', fmap_concat, output_shape_from=encoder['l2']['fmap'])\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl6']['fmap'], encoder['l2']['fmap']], axis=3)\n",
        "            decoder['dl7'] = self._build_decoder_layer('dl7', fmap_concat, output_shape_from=encoder['l1']['fmap'])\n",
        "\n",
        "            fmap_concat = tf.concat([decoder['dl7']['fmap'], encoder['l1']['fmap']], axis=3)\n",
        "            decoder['dl8'] = self._build_decoder_layer('dl8', fmap_concat, output_shape_from=self._inputs)\n",
        "\n",
        "            with tf.variable_scope('cl9'):\n",
        "                cl9 = dict()\n",
        "                cl9['filters'] = tf.get_variable('filters', [4, 4, get_shape(decoder['dl8']['fmap'])[-1], self._ochan])\n",
        "                cl9['conv'] =  tf.nn.conv2d(decoder['dl8']['fmap'], cl9['filters'], strides=[1, 1, 1, 1], padding='SAME')\n",
        "                cl9['fmap'] = tf.nn.tanh(cl9['conv'])\n",
        "                decoder['cl9'] = cl9\n",
        "        return decoder\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zH-TheQ_YRK_"
      },
      "source": [
        "## pytorch\n",
        "### generator\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, down=True, act=\"relu\", use_dropout=False):\n",
        "        super(Block, self).__init__()\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False, padding_mode=\"reflect\")\n",
        "            if down\n",
        "            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.ReLU() if act == \"relu\" else nn.LeakyReLU(0.2),\n",
        "        )\n",
        "\n",
        "        self.use_dropout = use_dropout\n",
        "        self.dropout = nn.Dropout(0.5)\n",
        "        self.down = down\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        return self.dropout(x) if self.use_dropout else x\n",
        "\n",
        "\n",
        "class P2P_Generator(nn.Module):\n",
        "    def __init__(self, in_channels=3, features=64):\n",
        "        super().__init__()\n",
        "        self.initial_down = nn.Sequential(\n",
        "            nn.Conv2d(in_channels, features, 4, 2, 1, padding_mode=\"reflect\"),\n",
        "            nn.LeakyReLU(0.2),\n",
        "        )\n",
        "        self.down1 = Block(features, features * 2, down=True, act=\"leaky\", use_dropout=False)\n",
        "        self.down2 = Block(\n",
        "            features * 2, features * 4, down=True, act=\"leaky\", use_dropout=False\n",
        "        )\n",
        "        self.down3 = Block(\n",
        "            features * 4, features * 8, down=True, act=\"leaky\", use_dropout=False\n",
        "        )\n",
        "        self.down4 = Block(\n",
        "            features * 8, features * 8, down=True, act=\"leaky\", use_dropout=False\n",
        "        )\n",
        "        self.down5 = Block(\n",
        "            features * 8, features * 8, down=True, act=\"leaky\", use_dropout=False\n",
        "        )\n",
        "        self.down6 = Block(\n",
        "            features * 8, features * 8, down=True, act=\"leaky\", use_dropout=False\n",
        "        )\n",
        "        self.bottleneck = nn.Sequential(\n",
        "            nn.Conv2d(features * 8, features * 8, 4, 2, 1), nn.ReLU()\n",
        "        )\n",
        "\n",
        "        self.up1 = Block(features * 8, features * 8, down=False, act=\"relu\", use_dropout=True)\n",
        "        self.up2 = Block(\n",
        "            features * 8 * 2, features * 8, down=False, act=\"relu\", use_dropout=True\n",
        "        )\n",
        "        self.up3 = Block(\n",
        "            features * 8 * 2, features * 8, down=False, act=\"relu\", use_dropout=True\n",
        "        )\n",
        "        self.up4 = Block(\n",
        "            features * 8 * 2, features * 8, down=False, act=\"relu\", use_dropout=False\n",
        "        )\n",
        "        self.up5 = Block(\n",
        "            features * 8 * 2, features * 4, down=False, act=\"relu\", use_dropout=False\n",
        "        )\n",
        "        self.up6 = Block(\n",
        "            features * 4 * 2, features * 2, down=False, act=\"relu\", use_dropout=False\n",
        "        )\n",
        "        self.up7 = Block(features * 2 * 2, features, down=False, act=\"relu\", use_dropout=False)\n",
        "        self.final_up = nn.Sequential(\n",
        "            nn.ConvTranspose2d(features * 2, in_channels, kernel_size=4, stride=2, padding=1),\n",
        "            nn.Tanh(),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        d1 = self.initial_down(x)\n",
        "        d2 = self.down1(d1)\n",
        "        d3 = self.down2(d2)\n",
        "        d4 = self.down3(d3)\n",
        "        d5 = self.down4(d4)\n",
        "        d6 = self.down5(d5)\n",
        "        d7 = self.down6(d6)\n",
        "        bottleneck = self.bottleneck(d7)\n",
        "        up1 = self.up1(bottleneck)\n",
        "        up2 = self.up2(torch.cat([up1, d7], 1))\n",
        "        up3 = self.up3(torch.cat([up2, d6], 1))\n",
        "        up4 = self.up4(torch.cat([up3, d5], 1))\n",
        "        up5 = self.up5(torch.cat([up4, d4], 1))\n",
        "        up6 = self.up6(torch.cat([up5, d3], 1))\n",
        "        up7 = self.up7(torch.cat([up6, d2], 1))\n",
        "        return self.final_up(torch.cat([up7, d1], 1))\n",
        "\n",
        "### discriminator\n",
        "\n",
        "class CNNBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, stride):\n",
        "        super(CNNBlock, self).__init__()\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_channels, out_channels, 4, stride, 1, bias=False, padding_mode=\"reflect\"\n",
        "            ),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.LeakyReLU(0.2),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.conv(x)\n",
        "\n",
        "\n",
        "class P2P_Discriminator(nn.Module):\n",
        "    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n",
        "        super().__init__()\n",
        "        self.initial = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_channels * 2,\n",
        "                features[0],\n",
        "                kernel_size=4,\n",
        "                stride=2,\n",
        "                padding=1,\n",
        "                padding_mode=\"reflect\",\n",
        "            ),\n",
        "            nn.LeakyReLU(0.2),\n",
        "        )\n",
        "\n",
        "        layers = []\n",
        "        in_channels = features[0]\n",
        "        for feature in features[1:]:\n",
        "            layers.append(\n",
        "                CNNBlock(in_channels, feature, stride=1 if feature == features[-1] else 2),\n",
        "            )\n",
        "            in_channels = feature\n",
        "\n",
        "        layers.append(\n",
        "            nn.Conv2d(\n",
        "                in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode=\"reflect\"\n",
        "            ),\n",
        "        )\n",
        "\n",
        "        self.model = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x, y):\n",
        "        x = torch.cat([x, y], dim=1)\n",
        "        x = self.initial(x)\n",
        "        x = self.model(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4wwM_-u9mMu0"
      },
      "source": [
        "## Training GAN models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2h7O6jv7DdU4"
      },
      "source": [
        "### Save Checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJs1IYzvDcvE"
      },
      "source": [
        "def save_checkpoint(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n",
        "    print(\"=> Saving checkpoint\")\n",
        "    checkpoint = {\n",
        "        \"state_dict\": model.state_dict(),\n",
        "        \"optimizer\": optimizer.state_dict(),\n",
        "    }\n",
        "    torch.save(checkpoint, filename)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UnbJ9LsStpM8"
      },
      "source": [
        "### Weights Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYKHqRgttNJN"
      },
      "source": [
        "def initialize_weights(model):\n",
        "    # Initializes weights according to the DCGAN paper\n",
        "    for m in model.modules():\n",
        "        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):\n",
        "            nn.init.normal_(m.weight.data, 0.0, 0.02)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M8nGXusgttf6"
      },
      "source": [
        "### Optimization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rF_lfVr3qtF5"
      },
      "source": [
        "##################################### select optimalization\n",
        "opt = ['Adam', 'RMS'] # 0, 1\n",
        "def Sel_Opt(opt, gen, disc):\n",
        "  if opt == 0:\n",
        "    opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
        "    opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
        "    return opt_gen, opt_disc\n",
        "  elif opt == 1:\n",
        "    opt_gen = optim.RMSprop(gen.parameters(), lr=LEARNING_RATE)\n",
        "    opt_critic = optim.RMSprop(disc.parameters(), lr=LEARNING_RATE)\n",
        "    return opt_gen, opt_critic "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g-CV7oNStwN1"
      },
      "source": [
        "### Loss function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NeW6vSpEt2CT"
      },
      "source": [
        "criterion =  nn.BCELoss()\n",
        "criterionP2P = nn.BCEWithLogitsLoss()\n",
        "l1_loss = nn.L1Loss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9GPX7mBVt3ka"
      },
      "source": [
        "### Train functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pVrRveEJ5B0h"
      },
      "source": [
        "##################################### train discriminator and generator\n",
        "def trainDisc(real,fake,disc, opt_disc):\n",
        "  print(fake.shape)\n",
        "  print(real.shape)\n",
        "  print(ok)\n",
        "  disc_real = disc(real).reshape(-1)\n",
        "  loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))\n",
        "  disc_fake = disc(fake.detach()).reshape(-1)\n",
        "  loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))\n",
        "  loss_disc = (loss_disc_real + loss_disc_fake) / 2\n",
        "  disc.zero_grad()\n",
        "  loss_disc.backward()\n",
        "  opt_disc.step()\n",
        "  return loss_disc\n",
        "\n",
        "def trainGen(fake, disc, gen, opt_gen):\n",
        "  output = disc(fake).reshape(-1)\n",
        "  loss_gen = criterion(output, torch.ones_like(output))\n",
        "  gen.zero_grad()\n",
        "  loss_gen.backward()\n",
        "  opt_gen.step()\n",
        "  return loss_gen\n",
        "\n",
        "def trainDisc_WGAN(real, CRITIC_ITERATIONS, BATCH_SIZE, gen, disc, opt_disc):\n",
        "  for _ in range(CRITIC_ITERATIONS):\n",
        "    noise = torch.randn(BATCH_SIZE, Z_DIM, 1, 1).to(device)\n",
        "    fake = gen(noise)\n",
        "    critic_real = disc(real).reshape(-1)\n",
        "    critic_fake = disc(fake).reshape(-1)\n",
        "    loss_critic = -(torch.mean(critic_real) - torch.mean(critic_fake))\n",
        "    disc.zero_grad()\n",
        "    loss_critic.backward(retain_graph=True)\n",
        "    opt_disc.step()\n",
        "\n",
        "    # clip critic weights between -0.01, 0.01\n",
        "    for p in disc.parameters():\n",
        "        p.data.clamp_(-WEIGHT_CLIP, WEIGHT_CLIP)\n",
        "    return loss_critic, fake\n",
        "\n",
        "def trainGen_WGAN(fake, gen, opt_gen, disc):\n",
        "  gen_fake = disc(fake).reshape(-1)\n",
        "  loss_gen = -torch.mean(gen_fake)\n",
        "  gen.zero_grad()\n",
        "  loss_gen.backward()\n",
        "  opt_gen.step()\n",
        "  return loss_gen\n",
        "\n",
        "def trainDisc_P2P(disc, gen, opt_disc, criterionP2P, d_scaler, real, y):\n",
        "  # Train Discriminator\n",
        "      with torch.cuda.amp.autocast():\n",
        "          y_fake = gen(real)\n",
        "          D_real = disc(real, y)\n",
        "          D_real_loss = criterionP2P(D_real, torch.ones_like(D_real))\n",
        "          D_fake = disc(real, y_fake.detach())\n",
        "          D_fake_loss = criterionP2P(D_fake, torch.zeros_like(D_fake))\n",
        "          D_loss = (D_real_loss + D_fake_loss) / 2\n",
        "\n",
        "      disc.zero_grad()\n",
        "      d_scaler.scale(D_loss).backward()\n",
        "      d_scaler.step(opt_disc)\n",
        "      d_scaler.update()\n",
        "      return disc, opt_disc, D_loss, y_fake\n",
        "\n",
        "\n",
        "\n",
        "def trainGen_P2P(disc, gen, opt_gen, l1_loss, criterionP2P, g_scaler, real, y, y_fake):\n",
        "      # # Train Discriminator\n",
        "      # with torch.cuda.amp.autocast():\n",
        "      #     y_fake = gen(real)\n",
        "      #     D_real = disc(real, y)\n",
        "      #     D_real_loss = bce(D_real, torch.ones_like(D_real))\n",
        "      #     D_fake = disc(real, y_fake.detach())\n",
        "      #     D_fake_loss = bce(D_fake, torch.zeros_like(D_fake))\n",
        "      #     D_loss = (D_real_loss + D_fake_loss) / 2\n",
        "\n",
        "      # disc.zero_grad()\n",
        "      # d_scaler.scale(D_loss).backward()\n",
        "      # d_scaler.step(opt_disc)\n",
        "      # d_scaler.update()\n",
        "\n",
        "      # Train generator\n",
        "      with torch.cuda.amp.autocast():\n",
        "          D_fake = disc(real, y_fake)\n",
        "          G_fake_loss = criterionP2P(D_fake, torch.ones_like(D_fake))\n",
        "          L1 = l1_loss(y_fake, y) * 100\n",
        "          G_loss = G_fake_loss + L1\n",
        "\n",
        "      opt_gen.zero_grad()\n",
        "      g_scaler.scale(G_loss).backward()\n",
        "      g_scaler.step(opt_gen)\n",
        "      g_scaler.update()\n",
        "      return gen, opt_gen, G_loss\n",
        "\n",
        "\n",
        "\n",
        "# def trainDisc_P2P():\n",
        "\n",
        "def tf_Train_DCdisc(loss_fn, batch_size, discriminator, fake, real, opt_disc):\n",
        "  with tf.GradientTape() as disc_tape:\n",
        "    print(real.shape)\n",
        "    real = tf.image.resize(real, [64, 64])\n",
        "    print(discriminator(real).shape)\n",
        "    print(fake.shape)\n",
        "    print((tf.ones((batch_size, 1))).shape)\n",
        "    loss_disc_real = loss_fn(tf.ones((batch_size, 1)), discriminator(real))\n",
        "    loss_disc_fake = loss_fn(tf.zeros(batch_size, 1), discriminator(fake))\n",
        "    loss_disc = (loss_disc_real + loss_disc_fake)/2\n",
        "\n",
        "    grads = disc_tape.gradient(loss_disc, discriminator.trainable_weights)\n",
        "    opt_disc.apply_gradients(\n",
        "        zip(grads, discriminator.trainable_weights)\n",
        "    )\n",
        "\n",
        "def tf_Train_DCgen(generator, discriminator, batch_size, loss_fn, random_latent_vectors, opt_gen):\n",
        "  with tf.GradientTape() as gen_tape:\n",
        "    fake = generator(random_latent_vectors)\n",
        "    output = discriminator(fake)\n",
        "    loss_gen = loss_fn(tf.ones(batch_size, 1), output)\n",
        "\n",
        "    grads = gen_tape.gradient(loss_gen, generator.trainable_weights)\n",
        "    opt_gen.apply_gradients(\n",
        "        zip(grads, generator.trainable_weights)\n",
        "    )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uQGvxgYot7Lh"
      },
      "source": [
        "### plot functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-5cXWmJsSFZB"
      },
      "source": [
        "################################# plot\n",
        "def pltLoss(epoch, dataloader, batch_idx, loss_disc, loss_gen, fixed_noise, gen, real, GANID):\n",
        "  print( f\"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(dataloader)} \\\n",
        "                  Loss D: {loss_disc:.4f}, loss G: {loss_gen:.4f}\")\n",
        "\n",
        "  with torch.no_grad():\n",
        "      if (GANID == 3) or(GANID == 1)or(GANID == 2):\n",
        "        fake = gen(fixed_noise)\n",
        "        # print()\n",
        "        # fake = gen(y).reshape(-1, 1, 28, 28)\n",
        "        # print(y.shape)\n",
        "        # print('================')\n",
        "        # print(real.shape)\n",
        "        # print(ok)\n",
        "      else:\n",
        "        fake = gen(fixed_noise).reshape(-1, 1, 28, 28)\n",
        "        # print(fake.shape)\n",
        "        # print(ok)\n",
        "        real = real.reshape(-1, 1, 28, 28)\n",
        "      # take out (up to) 32 examples\n",
        "      print('%%%%%%%%%%%%%%%', real.shape)\n",
        "      print('%%%%%%%%%%%%%%%%', fake.shape)\n",
        "      img_grid_real = torchvision.utils.make_grid(\n",
        "          real[:1], normalize=True\n",
        "      )\n",
        "      img_grid_fake = torchvision.utils.make_grid(\n",
        "          fake[:1], normalize=True\n",
        "      )\n",
        "      \n",
        "      plt.figure()\n",
        "      plt.xticks([])\n",
        "      plt.yticks([])\n",
        "      plt.grid(False)\n",
        "      print('img grid real--------------------',img_grid_real[0].shape)\n",
        "      imgreal = img_grid_fake[0].cpu().numpy()\n",
        "      print(imgreal*255)\n",
        "      plt.imshow(cv2.cvtColor(imgreal, cv2.COLOR_BGR2RGB))\n",
        "      plt.show()\n",
        "      imgfake = (img_grid_fake[0].cpu().numpy())*255\n",
        "      plt.imshow(cv2.cvtColor(imgfake, cv2.COLOR_BGR2RGB))\n",
        "      plt.show()\n",
        "      # print('$$$$$$$$$$$$', img_grid_real.shape)\n",
        "      # plt.figure(figsize=(10,10))\n",
        "      # # plt.imshow(img_grid_real)\n",
        "\n",
        "      # for i in range(25):\n",
        "      #     plt.subplot(5,5,i+1)\n",
        "      #     plt.xticks([])\n",
        "      #     plt.yticks([])\n",
        "      #     plt.grid(False)\n",
        "      #     print(img_grid_real[i].shape)\n",
        "      #     imgreal = img_grid_real[i].cpu().numpy()\n",
        "      #     print(imgreal.shape)\n",
        "      #     plt.imshow(cv2.cvtColor(imgreal, cv2.COLOR_BGR2RGB))\n",
        "\n",
        "      #     # The CIFAR labels happen to be arrays, \n",
        "      #     # which is why you need the extra index\n",
        "      # plt.show()\n",
        "      # plt.figure(figsize=(10,10))\n",
        "      # # plt.imshow(img_grid_real)\n",
        "      # for i in range(25):\n",
        "      #     plt.subplot(5,5,i+1)\n",
        "      #     plt.xticks([])\n",
        "      #     plt.yticks([])\n",
        "      #     plt.grid(False)\n",
        "      #     #plt.imshow(x_train[i])\n",
        "      #     imgfake = img_grid_fake[i].cpu().numpy()\n",
        "      #     plt.imshow(cv2.cvtColor(imgfake, cv2.COLOR_BGR2RGB))\n",
        "\n",
        "      #     # The CIFAR labels happen to be arrays, \n",
        "      #     # which is why you need the extra index\n",
        "      # plt.show()\n",
        "\n",
        "      # writer_real.add_image(\"Real\", img_grid_real, global_step=step)\n",
        "      # writer_fake.add_image(\"Fake\", img_grid_fake, global_step=step)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vjAuESmv-GDO"
      },
      "source": [
        "## Assemble training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iF7MesapA_xC"
      },
      "source": [
        "# dataset = keras.preprocessing.image_dataset_from_directory(\n",
        "#     directory=\"celeb_dataset\", label_mode=None, image_size=(64, 64), batch_size=32,\n",
        "#     shuffle=True\n",
        "# ).map(lambda x: x/255.0)\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WjLwbV31wjXz"
      },
      "source": [
        "def trainFunc(IMAGE_SIZE):\n",
        "  if usePt == 0:\n",
        "    if GANID == 1: # DCGAN\n",
        "      opt_gen = keras.optimizers.Adam(1e-4)\n",
        "      opt_disc = keras.optimizers.Adam(1e-4)\n",
        "      loss_fn = keras.losses.BinaryCrossentropy()\n",
        "      gen = generator_tf_dc\n",
        "      disc = discriminator_tf_dc\n",
        "    elif GANID == 2: # WGAN\n",
        "      gen = generator_tf_wgan(input, random_dim, is_train, reuse=False)\n",
        "      disc = discriminator_tf_wgan\n",
        "      opt_gen = \n",
        "      opt_disc = \n",
        "      loss_fn_g = \n",
        "      loss_fn_d = \n",
        "\n",
        "\n",
        "      with tf.variable_scope('input'):\n",
        "        # 模型中的输入数据部分\n",
        "        real_image = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 1], name='real_image')\n",
        "        random_input = tf.placeholder(tf.float32, shape=[None, Z_DIM], name='rand_input')\n",
        "        is_train = tf.placeholder(tf.bool, name='is_train')\n",
        "\n",
        "\n",
        "    for epoch in range(10):\n",
        "      for idx, (real, y_train) in enumerate(train_dataset):\n",
        "        batch_size = real.shape[0]\n",
        "        print(batch_size)\n",
        "        random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))\n",
        "        fake = gen(random_latent_vectors)\n",
        "        # if idx % 100 == 0:\n",
        "        #   img = keras.preprocessing.image.array_to_img(fake[0])\n",
        "        #   img.save(f\"generated_images/generated_img{epoch}_{idx}_.png\")\n",
        "\n",
        "        ### Train Discriminator: max log(D(x)) + log(1 - D(G(z))\n",
        "        tf_Train_DCdisc(loss_fn, batch_size, disc, fake, real, opt_disc)\n",
        "        ### Train Generator min log(1 - D(G(z)) <-> max log(D(G(z))\n",
        "        tf_Train_DCgen(gen, disc, batch_size, loss_fn, random_latent_vectors, opt_gen)\n",
        "        print('finshied------------------')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "  else:\n",
        "    if GANID == 0:\n",
        "      IMAGE_SIZE = 28 * 28 * 1\n",
        "      disc = SimpleDiscriminator(IMAGE_SIZE).to(device)\n",
        "      gen  = SimpleGenerator(Z_DIM, IMAGE_SIZE).to(device)\n",
        "      fixed_noise = torch.randn(BATCH_SIZE, Z_DIM).to(device)\n",
        "    elif GANID == 1:\n",
        "      IMAGE_SIZE = IMAGE_SIZE\n",
        "      gen = DCGAN_Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)\n",
        "      disc = DCGAN_Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)\n",
        "      fixed_noise = torch.randn(BATCH_SIZE, Z_DIM, 1, 1).to(device)\n",
        "    elif GANID == 2:\n",
        "      gen = WGAN_Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)\n",
        "      disc = WGAN_Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)\n",
        "      fixed_noise = torch.randn(BATCH_SIZE, Z_DIM, 1, 1).to(device)\n",
        "    elif GANID == 3:\n",
        "      gen = P2P_Generator(in_channels=3, features=64).to(device)\n",
        "      disc = P2P_Discriminator(in_channels=3).to(device)\n",
        "      # fixed_noise = torch.randn(BATCH_SIZE, Z_DIM, 1, 1).to(device)\n",
        "      \n",
        "      \n",
        "    initialize_weights(gen)\n",
        "    initialize_weights(disc)\n",
        "    dataset = datasetReturn(GANID)\n",
        "    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
        "    if GANID == 3:\n",
        "      dataset = MapDataset(root_dir='/content/drive/MyDrive/Colab Notebooks/datasets/maps/train/')\n",
        "      dataloader = DataLoader(\n",
        "            dataset,\n",
        "            batch_size=BATCH_SIZE,\n",
        "            shuffle=True,\n",
        "        )\n",
        "    opt_gen, opt_disc = Sel_Opt(0, gen, disc) # optimal function selection\n",
        "\n",
        "    \n",
        "    writer_real = SummaryWriter(f\"logs/real\")\n",
        "    writer_fake = SummaryWriter(f\"logs/fake\")\n",
        "    step = 0\n",
        "\n",
        "    gen.train()\n",
        "    disc.train()\n",
        "    ################################################# train start\n",
        "    for epoch in range(NUM_EPOCHS):\n",
        "      # Target labels not needed! <3 unsupervised\n",
        "      for batch_idx, (real, y) in enumerate(dataloader):\n",
        "        real = real.to(device)\n",
        "        y = y.to(device)\n",
        "        # print(real)\n",
        "        # print(fake1)\n",
        "        # print(ok)\n",
        "        # print('real',real.shape)\n",
        "        if GANID == 0:\n",
        "          real = real.view(-1, 784).to(device)\n",
        "          noise = torch.randn(batch_size, Z_DIM).to(device)\n",
        "          # print('noise',noise.shape)\n",
        "          fake = gen(noise)\n",
        "          # print('',fake.shape)\n",
        "        elif (GANID == 1):\n",
        "          noise = torch.randn(BATCH_SIZE, Z_DIM, 1, 1).to(device)\n",
        "          fake = gen(noise)\n",
        "          print(real.shape)\n",
        "        elif (GANID == 3):\n",
        "          g_scaler = torch.cuda.amp.GradScaler()\n",
        "          d_scaler = torch.cuda.amp.GradScaler()\n",
        "          \n",
        "\n",
        "        if GANID == 2:\n",
        "          loss_disc, fake = trainDisc_WGAN(real, CRITIC_ITERATIONS, BATCH_SIZE, gen, disc, opt_disc)\n",
        "          loss_gen = trainGen_WGAN(fake, gen, opt_gen, disc)\n",
        "        elif (GANID == 0)or(GANID == 1):\n",
        "          loss_disc = trainDisc(real, fake, disc, opt_disc)\n",
        "          loss_gen = trainGen(fake, disc, gen, opt_gen)\n",
        "        elif (GANID == 3):\n",
        "          disc, opt_disc, loss_disc, y_fake = trainDisc_P2P(disc, gen, opt_disc, criterionP2P, d_scaler, real, y)\n",
        "          gen, opt_gen, loss_gen = trainGen_P2P(disc, gen, opt_gen, l1_loss, criterionP2P, g_scaler, real, y, y_fake)\n",
        "          fixed_noise = y\n",
        "\n",
        "        # Print losses occasionally and print to tensorboard\n",
        "        if batch_idx % 100 == 0:\n",
        "          pltLoss(epoch, dataloader, batch_idx, loss_disc, loss_gen, fixed_noise, gen, real, GANID)\n",
        "          step += 1\n",
        "      if (epoch % 100 == 0) or (epoch == NUM_EPOCHS-1):\n",
        "        save_checkpoint(gen, opt_gen, filename=\"gen_checkpoint_\"+str(epoch) +\".pth.tar\")\n",
        "        save_checkpoint(disc, opt_disc, filename = \"disc_checkpoint_\"+str(epoch) +\".pth.tar\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Nwwb-q9Tkl8"
      },
      "source": [
        "trainFunc(IMAGE_SIZE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jeRdRNJF2kwU"
      },
      "source": [
        "# More resources:\n",
        "\n",
        "* 3Blue1Brown Neural Network [video tutorials](https://www.youtube.com/playlist?list=PLZHQObOWTQDNU6R1_67000Dx_ZCJB-3pi) \n",
        "* Deep Learning Video Lectures by Prof. Andreas Maier [Winter 20/21](https://www.youtube.com/watch?v=SCFToE1vM2U&list=PLpOGQvPCDQzvJEPFUQ3mJz72GJ95jyZTh)\n",
        "* Some of the code in this notebook is taken from [here](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch)\n",
        "* Calculating number of parameters in [CNN](https://towardsdatascience.com/understanding-and-calculating-the-number-of-parameters-in-convolution-neural-networks-cnns-fc88790d530d)\n",
        "* Some of the code in this notebook is taken from [here](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/images/cnn.ipynb)\n",
        "* https://imerit.net/blog/top-13-machine-learning-image-classification-datasets-all-pbm/\n",
        "* https://nihcc.app.box.com/v/ChestXray-NIHCC\n",
        "* https://www.kaggle.com/xhlulu/recursion-cellular-image-classification-224-jpg\n",
        "* https://www.tensorflow.org/datasets/catalog/patch_camelyon\n",
        "* https://www.youtube.com/playlist?list=PLhhyoLH6IjfwIp8bZnzX8QR30TRcHO8Va\n",
        "* https://developers.google.com/machine-learning/gan/gan_structure\n",
        "* https://machinelearningmastery.com/impressive-applications-of-generative-adversarial-networks/"
      ]
    }
  ]
}