evaluate_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "evaluate GAN.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "nPIXJBUPMuTr",
        "DOKYXZp3Bbnb",
        "l-vj_z3Pleut"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/buoi/conditional-face-GAN/blob/main/evaluate_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": "gMBqo7it5ghH"
      },
      "source": [
        "# ⚙️ Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G3sABQInUGmh"
      },
      "source": [
        "## Fix random seeds"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z9Kyg2cXUFnd"
      },
      "source": [
        "SEED = 11\n",
        "import os\n",
        "os.environ['PYTHONHASHSEED']=str(SEED)\n",
        "import random\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "random.seed(SEED)\n",
        "np.random.seed(SEED)\n",
        "tf.random.set_seed(SEED)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ikEHqqiEUNwP"
      },
      "source": [
        "## Automatic Mixed Precision\n",
        "if supported by GPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "emWtF7RxfhPb"
      },
      "source": [
        "_, gpu_name  = !nvidia-smi --query-gpu=gpu_name --format=csv\n",
        "\n",
        "if gpu_name == 'Tesla T4':\n",
        "    from tensorflow.keras import mixed_precision\n",
        "\n",
        "    policy = mixed_precision.Policy('mixed_float16')\n",
        "    mixed_precision.set_global_policy(policy)\n",
        "\n",
        "!nvidia-smi -L"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UOswoZB0oeA2"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-RSWEMo35ghH"
      },
      "source": [
        "from scipy import linalg\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from keras.applications.inception_v3 import InceptionV3, preprocess_input\n",
        "\n",
        "from tensorflow.keras.layers import Layer, Input, Dense, Reshape, Flatten\n",
        "from tensorflow.keras.layers import Conv2D, Conv2DTranspose, ReLU, LeakyReLU\n",
        "from tensorflow.keras.layers import Dropout, Embedding, Concatenate, Add, Activation\n",
        "from tensorflow.keras.layers import GlobalAveragePooling2D, UpSampling2D, BatchNormalization\n",
        "import tensorflow.keras.backend as K\n",
        "\n",
        "from tensorflow.python.keras.utils import conv_utils\n",
        "from tensorflow.keras.initializers import RandomNormal\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "\n",
        "!pip install tensorflow_addons\n",
        "import tensorflow_addons as tfa\n",
        "from tensorflow_addons.layers import SpectralNormalization\n",
        "\n",
        "import gdown\n",
        "from zipfile import ZipFile\n",
        "\n",
        "from tqdm.notebook import tqdm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bL758kqufZWu"
      },
      "source": [
        "## Download CelebA\n",
        "\n",
        "We'll use face images from the CelebA dataset, resized to 64x64."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "40CS8jHvfZWu"
      },
      "source": [
        "try:\n",
        "    os.makedirs(\"celeba_gan\")\n",
        "    url = \"https://drive.google.com/uc?id=1O7m1010EJjLE5QxLZiM9Fpjs7Oj6e684\"\n",
        "    output = \"celeba_gan/img_align_celeba.zip\"\n",
        "    gdown.download(url, output, quiet=True)\n",
        "\n",
        "    with ZipFile(\"celeba_gan/img_align_celeba.zip\", \"r\") as zipobj:\n",
        "        zipobj.extractall(\"celeba_gan\")\n",
        "        \n",
        "except FileExistsError:\n",
        "    print(\"Dataset Already downloaded\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oZ32Qan0YOBX"
      },
      "source": [
        "#Download labels from public github, they have been processed in a 0,1 csv file\n",
        "!wget -q -O \"/content/celeba_gan/list_attr_celeba01.csv.zip\" \"https://github.com/buoi/conditional-face-GAN/blob/main/list_attr_celeba01.csv.zip?raw=true\" \n",
        "!unzip -o \"/content/celeba_gan/list_attr_celeba01.csv.zip\" -d \"/content/celeba_gan\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_uR3O5gxul1x"
      },
      "source": [
        "## Dataset preprocessing functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k5D5SOGWhsNc"
      },
      "source": [
        "# image utils functions\n",
        "\n",
        "def conv_range(in_range=(-1,1), out_range=(0,255)):\n",
        "    \"\"\" Returns range conversion function\"\"\"\n",
        "\n",
        "    # compute means and spans once\n",
        "    in_mean, out_mean = np.mean(in_range), np.mean(out_range)\n",
        "    in_span, out_span = np.ptp(in_range), np.ptp(out_range)\n",
        "\n",
        "    # return function\n",
        "    def convert_img_range(in_img):\n",
        "        out_img = (in_img - in_mean) / in_span\n",
        "        out_img = out_img * out_span + out_mean\n",
        "        return out_img\n",
        "\n",
        "    return convert_img_range\n",
        "\n",
        "def crop128(img):\n",
        "    #return img[:, 77:141, 57:121]# 64,64 center crop\n",
        "    return img[:, 45:173, 25:153] # 128,128 center crop\n",
        "\n",
        "def resize64(img):\n",
        "    return tf.image.resize(img, (64,64), antialias=True, method='bilinear')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2XeNjuoHAiDT"
      },
      "source": [
        "## FID utilities"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6hOVFi5r_Fqu"
      },
      "source": [
        "#fid real images embeddings\n",
        "inception_model = InceptionV3(include_top=False, \n",
        "                              weights=\"imagenet\", \n",
        "                              pooling='avg',\n",
        "                              input_shape=(299, 299, 3))\n",
        "inception_model.trainable = False\n",
        "\n",
        "\n",
        "def define_real_prep():\n",
        "    img = layers.Input(shape=IMAGE_SHAPE, name='img_input')\n",
        "    img_up = tf.image.resize(img, (299, 299), method='bilinear', antialias=True, name='upsample_bilinear')\n",
        "    output = inception_model(img_up, training=False)\n",
        "    real_prep = tf.keras.Model(inputs=[img], outputs=[output])\n",
        "\n",
        "    return real_prep\n",
        "\n",
        "def generate_dataset_embeddings(dataset, fid_count=10000):\n",
        "\n",
        "    #tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)\n",
        "    real_prep = define_real_prep()\n",
        "    #tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.WARN)\n",
        "\n",
        "    real_embeddings = np.zeros((FID_COUNT, 2048), dtype='float32')\n",
        "    real_attr = np.zeros((FID_COUNT, N_ATTRIBUTES), dtype='float32')\n",
        "\n",
        "    fid_bs = 200\n",
        "    assert fid_count%fid_bs ==0, 'bad fid batch size or fid_count'\n",
        "    dataset = dataset.unbatch().batch(fid_bs)\n",
        "    dataset.shuffle(210000)\n",
        "\n",
        "    for i in tqdm(range(fid_count//fid_bs)):\n",
        "        if N_ATTRIBUTES !=0:\n",
        "            batch_img, batch_attr = next(iter(dataset))\n",
        "            real_attr[i*fid_bs:(i+1)*fid_bs] = batch_attr\n",
        "        else:\n",
        "            batch_img = next(iter(dataset))\n",
        "\n",
        "        batch_embeddings = real_prep(batch_img)\n",
        "        real_embeddings[i*fid_bs:(i+1)*fid_bs] = batch_embeddings\n",
        "\n",
        "    return real_embeddings, real_attr\n",
        "\n",
        "\n",
        "def generate_fid_embeddings(model, real_attr, fid_count=10000, use_attributes=False):\n",
        "\n",
        "\n",
        "    #tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)\n",
        "    if N_ATTRIBUTES !=0 and use_attributes:\n",
        "        multi_evaluator = define_evaluator(model.generator, attribute_classifier)\n",
        "    else:\n",
        "        multi_evaluator = define_evaluator(model.generator) #1\n",
        "    #tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.WARN)\n",
        "    \n",
        "    gen_size = 100\n",
        "    assert fid_count % gen_size == 0, \"wrong gen_size or fid_count\"\n",
        "\n",
        "    generated_embeddings = np.zeros((fid_count, 2048), dtype='float32') #2\n",
        "    pred_attr = np.zeros((fid_count, N_ATTRIBUTES), dtype='float32')\n",
        "\n",
        "    for i in tqdm(range(fid_count//gen_size)):\n",
        "        if N_ATTRIBUTES !=0 and use_attributes:\n",
        "            batch_embeddings, batch_attr = multi_evaluator((\n",
        "                tf.random.normal([gen_size, LATENT_DIM]),\n",
        "                real_attr[i*gen_size:(i+1)*gen_size]))\n",
        "            pred_attr[i*gen_size:(i+1)*gen_size] = batch_attr\n",
        "            \n",
        "        elif N_ATTRIBUTES !=0: # call model with attributes but returns embeddings only #3\n",
        "            batch_embeddings = multi_evaluator((\n",
        "                tf.random.normal([gen_size, LATENT_DIM]),\n",
        "                real_attr[i*gen_size:(i+1)*gen_size]))\n",
        "        else:\n",
        "            batch_embeddings = multi_evaluator(\n",
        "                tf.random.normal([gen_size, LATENT_DIM]))\n",
        "            \n",
        "        generated_embeddings[i*gen_size:(i+1)*gen_size] = batch_embeddings\n",
        "\n",
        "    return generated_embeddings, pred_attr\n",
        "\n",
        "def define_evaluator(generator, attr_classifier=None):\n",
        "    \n",
        "    #z_input = layers.Input(shape=(LATENT_DIM), name='z_input')\n",
        "    #attr_input = layers.Input(shape=(N_ATTRIBUTES), name='attr_input')\n",
        "    #img = generator((z_input,attr_input))\n",
        "    img = generator.output\n",
        "\n",
        "    img_up299 = tf.image.resize(img, (299, 299), method='bilinear', antialias=True, name='upsample_bilinear')\n",
        "    output_repr = inception_model(img_up299, training=False,)\n",
        "\n",
        "    if N_ATTRIBUTES !=0 and attr_classifier is not None:\n",
        "        print(\"using classifier\")\n",
        "        attributes = attr_classifier(img) # mobilenet model does resizing.. does *2-1 also...\n",
        "        multi_evaluator = tf.keras.Model(inputs=generator.input, outputs=[output_repr, attributes])\n",
        "    else:\n",
        "        multi_evaluator = tf.keras.Model(inputs=generator.input, outputs=output_repr)\n",
        "\n",
        "    return multi_evaluator\n",
        "\n",
        "\"\"\"def define_evaluator(generator, attr_classifier=None):\n",
        "    \n",
        "    img = generator.output\n",
        "\n",
        "    img_up229 = tf.image.resize(img, (299, 299), method='bilinear', antialias=True, name='upsample_bilinear')\n",
        "    output_repr = inception_model(img_up229, training=False,)\n",
        "\n",
        "    if N_ATTRIBUTES !=0 and attr_classifier is not None:\n",
        "        print(\"using classifier\")\n",
        "        attributes = attr_classifier(img) # mobilenet model does resizing.. does *2-1 also...\n",
        "        multi_evaluator = tf.keras.Model(inputs=generator.inputs, outputs=[output_repr, attributes])\n",
        "    else:\n",
        "        multi_evaluator = tf.keras.Model(inputs=generator.inputs, outputs=output_repr)\n",
        "\n",
        "    return multi_evaluator\"\"\"\n",
        "\n",
        "\n",
        "def compute_fid(real_embeddings, generated_embeddings):\n",
        "    # compute mean and covariance statistics\n",
        "    print(\"computing FID on:\",real_embeddings.shape, generated_embeddings.shape)\n",
        "    mu1, sigma1 = real_embeddings.mean(axis=0), np.cov(real_embeddings, rowvar=False)\n",
        "    mu2, sigma2 = generated_embeddings.mean(axis=0), np.cov(generated_embeddings, rowvar=False)\n",
        "    # compute sum squared difference between means\n",
        "    ssdiff = np.sum((mu1 - mu2)**2.0)\n",
        "    # compute sqrt of product between cov\n",
        "    covmean = linalg.sqrtm(sigma1.dot(sigma2))\n",
        "    # check and correct imaginary numbers from sqrt\n",
        "    if np.iscomplexobj(covmean):\n",
        "      covmean = covmean.real\n",
        "    # compute fid score\n",
        "    trace = np.trace(sigma1 + sigma2 - 2.0 * covmean)\n",
        "    fid = ssdiff + trace\n",
        "    print(\"fid:\",fid, \"ssdiff:\",ssdiff, \"trace:\",trace)\n",
        "    return fid"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYtB7-eejd7X"
      },
      "source": [
        "# 📉 Evaluate model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "71HViZzUAm-7"
      },
      "source": [
        "## Load trained GAN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n4V7EUGfupKA"
      },
      "source": [
        "from collections import namedtuple\n",
        "ModelEntry = namedtuple('ModelEntry', '''entity, resume_id, model_name, exp_name, best_epoch, expected_fid, expected_f1, expected_acc, epoch_range, exp_avg_fid, exp_avg_f1, exp_avg_acc''')\n",
        "\n",
        "dcgan = ModelEntry('buio','t1qyadzp', 'dcgan','dcgan',\n",
        "                   'v24',25.64, 0, 0, (16,28), 24.29, 0,0)\n",
        "\n",
        "acgan2 = ModelEntry('buio','rn8xslip','acgan2_BNstdev','acgan2_BNstdev', \n",
        "                    'v22', 29.05, 0.9182, 0.9295, (20,31), 24.79, 0.918, 0.926)\n",
        "\n",
        "acgan10 = ModelEntry('buio','3ja6uvac','acgan10_nonseparBNstdev_split','acgan10_nonseparBNstdev_split',\n",
        "                     'v24', 26.89, 0.859, 0.785, (18,30), 24.59, 0.789,0.858)\n",
        "\n",
        "acgan40 = ModelEntry('buio','2ev65fpt','acgan40_BNstdev','acgan40_BNstdev',\n",
        "                     'v15', 28.23, 0.430, 0.842, (15,25), 27.72, 04.6, 0.851)\n",
        "\n",
        "\n",
        "acgan2_hd = ModelEntry('buio','6km6fdgr','acgan2_BNstdev_218x178','acgan2_BNstdev_218x178',\n",
        "                     'v11', 0,0,0 , 0, 0, 0)\n",
        "\n",
        "acgan10_hd = ModelEntry('buio','3v366skw','acgan40_BNstdev_218x178','acgan40_BNstdev_218x178',\n",
        "                     'v14', 0,0,0, 0, 0, 0)\n",
        "\n",
        "acgan40_hd = ModelEntry('buio','booicugb','acgan10_nonseparBNstdev_split_299_218x178','acgan10_nonseparBNstdev_split_299_218x178',\n",
        "                     'v14', 52.9, 0.410, 0.834, (12,15), 0, 0, 0)\n",
        "\n",
        "\n",
        "\n",
        "#1cr1a5w4 SAGAN_3 v31 buianifolli\n",
        "#2o3z6bqb SAGAN_5 v17 buianifolli\n",
        "#zscel8bz SAGAN_6 v29 buianifolli\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1FxRLbu6C1-D"
      },
      "source": [
        "model = acgan40_hd\n",
        "entity = model.entity\n",
        "project_name = \"GAN\" \n",
        "\n",
        "resume_id = model.resume_id\n",
        "\n",
        "model_name = model.model_name\n",
        "experiment_name = model.exp_name\n",
        "EPOCH_RANGE = model.epoch_range\n",
        "version = model.best_epoch\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IbREbwDIAcXV"
      },
      "source": [
        "# loads a model from Wandb\n",
        "IMAGE_RANGE='11'\n",
        "\n",
        "!pip install wandb > /dev/null\n",
        "!wandb login --relogin\n",
        "import wandb\n",
        "\n",
        "run = wandb.init(project=project_name, \n",
        "                  entity = entity,\n",
        "                  name=experiment_name, \n",
        "                  resume=resume_id)\n",
        "\n",
        "\n",
        "run_name = entity+'/GAN/'+ model_name + experiment_name\n",
        "artifact_run = run_name +':'+version\n",
        "\n",
        "artifact = run.use_artifact(artifact_run, type='model')\n",
        "artifact_dir = artifact.download()\n",
        "gan = keras.models.load_model(os.path.join(artifact_dir, model_name + experiment_name))\n",
        "\n",
        "\n",
        "IMAGE_SIZE = gan.discriminator.input_shape[1]\n",
        "if IMAGE_SIZE == 64:\n",
        "    IMAGE_SHAPE = (64,64,3)\n",
        "elif IMAGE_SIZE == 218:\n",
        "    IMAGE_SHAPE = (218,178,3)\n",
        "\n",
        "try\n",
        "    LATENT_DIM = gan.generator.input_shape[0][1]\n",
        "    N_ATTRIBUTES = gan.generator.input_shape[1][1]\n",
        "except TypeError:\n",
        "    LATENT_DIM = gan.generator.input_shape[1]\n",
        "    N_ATTRIBUTES =0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfErOFepg7sF"
      },
      "source": [
        "## 💾 Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0WO2104NnsGh"
      },
      "source": [
        "#@title Select Attributes {form-width: \"50%\", display-mode: \"both\" }\n",
        "\n",
        "#NUMBER_OF_ATTRIBUTES = \"10\" #@param [0, 2, 10, 12, 40]\n",
        "#N_ATTRIBUTES = int(NUMBER_OF_ATTRIBUTES)\n",
        "\n",
        "IMAGE_RANGE = '11'\n",
        "\n",
        "BATCH_SIZE =  64 #@param {type: \"number\"}\n",
        "if N_ATTRIBUTES == 2:\n",
        "    LABELS = [\"Male\", \"Smiling\"]\n",
        "\n",
        "elif N_ATTRIBUTES == 10:\n",
        "    LABELS = [\n",
        "          \"Mouth_Slightly_Open\", \"Wearing_Lipstick\", \"High_Cheekbones\", \"Male\", \"Smiling\", \n",
        "          \"Heavy_Makeup\", \"Wavy_Hair\", \"Oval_Face\", \"Pointy_Nose\", \"Arched_Eyebrows\"]\n",
        "\n",
        "elif N_ATTRIBUTES == 12:\n",
        "    LABELS = ['Wearing_Lipstick','Mouth_Slightly_Open','Male','Smiling',\n",
        "              'High_Cheekbones','Heavy_Makeup','Attractive','Young',\n",
        "              'No_Beard','Black_Hair','Arched_Eyebrows','Big_Nose']\n",
        "elif N_ATTRIBUTES == 40:\n",
        "    LABELS = [\n",
        "            '5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive',\n",
        "            'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose',\n",
        "            'Black_Hair', 'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows',\n",
        "            'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee', 'Gray_Hair',\n",
        "            'Heavy_Makeup', 'High_Cheekbones', 'Male', 'Mouth_Slightly_Open',\n",
        "            'Mustache', 'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin',\n",
        "            'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks', 'Sideburns',\n",
        "            'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings',\n",
        "            'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace',\n",
        "            'Wearing_Necktie', 'Young']\n",
        "\n",
        "else:\n",
        "    LABELS = [\"Male\", \"Smiling\"]# just for dataset creation\n",
        " \n",
        "\n",
        "# Take labels and a list of image locations in memory\n",
        "df = pd.read_csv(r\"/content/celeba_gan/list_attr_celeba01.csv\")\n",
        "im_list = df['image_id'].tolist()\n",
        "\n",
        "\n",
        "# load image at exact resolution\n",
        "dataset_train = keras.preprocessing.image_dataset_from_directory(\n",
        "    \"celeba_gan\", label_mode=\"int\", \n",
        "    labels = df[LABELS].values.tolist(),\n",
        "    image_size=(218, 178), \n",
        "    batch_size=BATCH_SIZE,\n",
        "    seed=SEED)\n",
        "\n",
        "\n",
        "def preprocess(dataset):\n",
        "    # extract the center crop\n",
        "    if IMAGE_SIZE == 64:\n",
        "        dataset = dataset.map(lambda x,y: (crop128(x),y) )\n",
        "        # resize\n",
        "        dataset = dataset.map(lambda x,y: (resize64(x),y) )\n",
        "    # convert image range\n",
        "    dataset = dataset.map(lambda x,y: (conv_range([0,255], [-1,1])(x),y) )\n",
        "    return dataset\n",
        "\n",
        "dataset_train = preprocess(dataset_train)\n",
        "\n",
        "if N_ATTRIBUTES !=0: \n",
        "    batch_img, batch_labels = next(iter(dataset_train))\n",
        "    #plt.title('\\n'.join([(1-label)*'Not_'+LABELS[i]+' •'*int(label) for i,label in enumerate(batch_labels[0])]))\n",
        "    plt.title('\\n'.join([' ☑  '+label*' '.join(LABELS[i].split('_')) for i,label in enumerate(batch_labels[0])]),loc='left',fontsize=15)\n",
        "          \n",
        "\n",
        "else: # remove labels\n",
        "    dataset_train = dataset_train.map(lambda x,y: x)\n",
        "    batch_img = next(iter(dataset_train))\n",
        "    \n",
        "plt.imshow(batch_img[0].numpy()/2+0.5)\n",
        "\n",
        "print(\"image in range: \", np.min(batch_img), np.max(batch_img))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hUi_ouNiRc4e"
      },
      "source": [
        "# Evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3afRYQMLlKL-"
      },
      "source": [
        "## Multi epochs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xwQGqYsoQ24N"
      },
      "source": [
        "EPOCH_RANGE = range(*model.epoch_range)\n",
        "FID_COUNT = 10000\n",
        "\n",
        "\n",
        "models = []\n",
        "for i in EPOCH_RANGE:\n",
        "\n",
        "    artifact_run = run_name +':v'+str(i)\n",
        "\n",
        "    artifact = run.use_artifact(artifact_run, type='model')\n",
        "    artifact_dir = artifact.download()\n",
        "    gan = keras.models.load_model(os.path.join(artifact_dir, model_name + experiment_name))\n",
        "    models.append(gan)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d_7KoEYPRnYq"
      },
      "source": [
        "#@title Load Multilabel Classifier Model{form-width: \"35%\", display-mode: \"both\" }\n",
        "\n",
        "def f1_metric_logits(y_true, y_pred):\n",
        "    y_pred = K.round(keras.activations.sigmoid(y_pred))\n",
        "    tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)\n",
        "    tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)\n",
        "    fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)\n",
        "    fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)\n",
        "\n",
        "    p = tp / (tp + fp + K.epsilon())\n",
        "    r = tp / (tp + fn + K.epsilon())\n",
        "\n",
        "    f1 = 2*p*r / (p+r+K.epsilon())\n",
        "    f1 = tf.where(tf.math.is_nan(f1), tf.zeros_like(f1), f1)\n",
        "    return K.mean(f1)\n",
        "\n",
        "\n",
        "def weighted_bce_from_logits(y_true, y_pred):\n",
        "\n",
        "    #inverted_frequency weighting\n",
        "    pos_weight= (1-attribute_frequency)*2\n",
        "    y_true = tf.cast(y_true,'float32')\n",
        "    bce = tf.nn.weighted_cross_entropy_with_logits(y_true, y_pred, pos_weight=pos_weight)\n",
        "    #weighted_bce = attribute_frequency*(1-y_true) *bce\n",
        "    #print(bce, weighted_bce)\n",
        "    return K.mean(bce,axis=-1)\n",
        "\n",
        "\n",
        "custom_objects = {'f1_metric_logits': f1_metric_logits,\n",
        "                  'weighted_bce_from_logits': weighted_bce_from_logits}\n",
        "\n",
        "if IMAGE_SIZE == 64:\n",
        "    if N_ATTRIBUTES == 2:\n",
        "        resume_id = \"2ckbukxu\"\n",
        "        model_name = experiment_name = \"mobilenet_2attr_weighted\"\n",
        "        \n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDI1ODc=/54e0c2174d2d72b2aad99d6f3edb9478\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDI1ODc=/c6ed38f84d49a73fcccc7cc3947d6976\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDI1ODc=/169bada1896162b8600275781a507077\"\n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDI1ODc=/014e7165dbc24bc0ed55e088bf56ff65\"\n",
        "\n",
        "\n",
        "    elif N_ATTRIBUTES == 10:\n",
        "        resume_id = \"1ciqd5nw\"\n",
        "        model_name = experiment_name = \"mobilenet_10attr_weighted\"\n",
        "\n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDkwODM=/b0843a4e7f983e156eefda7e218fac84\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDkwODM=/73898d1b00855e233b5eb9ad2f19edd6\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDkwODM=/7ef193295b72d42e6990c66d0e44c98a\" \n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0NDkwODM=/2ef2d12ecc8f79f17d55d430a185b8d5\"\n",
        "\n",
        "    elif N_ATTRIBUTES == 40:\n",
        "        resume_id = \"123ijj97\"\n",
        "        model_name = experiment_name = \"mobilenet_40attr_weighted\"\n",
        "        \n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0MzU4NzY=/c2696523ad6bc0a5c3e80eaa528158a0\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0MzU4NzY=/df7357b4fb17005995cd9c8ac37f1aa1\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0MzU4NzY=/ed6bb693bb3789ebae0e79ec2d8331a4\" \n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI0MzU4NzY=/c72922c943050df965eff6484dbfc954\"\n",
        "\n",
        "    else:\n",
        "        print(\"N_ATTRIBUTES=\", N_ATTRIBUTES, \"skipping classifier load\")\n",
        "        ENABLE_LOAD = False\n",
        "\n",
        "elif IMAGE_SIZE == 218:\n",
        "    if N_ATTRIBUTES == 2:\n",
        "        resume_id = '3b15ld8p'\n",
        "        model_name = experiment_name = \"mobilenet_2attr_218x178_weighted\"\n",
        "\n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1NzgwNzE=/f68713bc74adac9c3d8b4aee9064f76a\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1NzgwNzE=/65f63abdb95dbf59c398ab0a9314609b\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1NzgwNzE=/b313ebd8a28edf7b9fcc291f5d350f83\"\n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1NzgwNzE=/728b69786f7a1fd77cb0d6edb5b0b6d1\"\n",
        "\n",
        "    elif N_ATTRIBUTES == 10:\n",
        "        resume_id = \"3n56eda3\"\n",
        "        model_name = experiment_name = \"mobilenet_10attr_218x178_weighted\"\n",
        "\n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTkyMTk=/55a5f2dfe27eeb1873e05e8b7a4be86a\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTkyMTk=/a13318bc9a500bcde0cb020182dec064\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTkyMTk=/69be9cba2020b21eeb4bcb6bbf887136\"\n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTkyMTk=/65079a2e8655615837abbf7390aec2fc\"\n",
        "\n",
        "    elif N_ATTRIBUTES == 40:\n",
        "        resume_id = \"zd9xh3kq\"\n",
        "        model_name = experiment_name = \"mobilenet_40attr_218x178_weighted\"\n",
        "\n",
        "        keras_metadata_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTAyMjU=/58e2dc40dc2afb2ac5c771acc1d2dd78\"\n",
        "        saved_model_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTAyMjU=/22cb81a3b8ec24619f8204d383849c81\"\n",
        "        variables_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTAyMjU=/c80537ac89fa927fd68bcb6cd752a8a6\"\n",
        "        index_url = \"https://api.wandb.ai/artifactsV2/gcp-us/buio/QXJ0aWZhY3Q6MzI1OTAyMjU=/79ec8ead744b1bc4e19c098de9ff0f22\"\n",
        "\n",
        "    else:\n",
        "        print(\"N_ATTRIBUTES=\", N_ATTRIBUTES, \"skipping classifier load\")\n",
        "        ENABLE_LOAD = False\n",
        "\n",
        "ENABLE_WANDB = False\n",
        "if N_ATTRIBUTES !=0 :\n",
        "    if ENABLE_WANDB:\n",
        "        # load classifier artifact\n",
        "        project_name = \"GAN\"\n",
        "        !pip install wandb > /dev/null\n",
        "        !wandb login\n",
        "        import wandb\n",
        "        from wandb.keras import WandbCallback\n",
        "        run = wandb.init(project=project_name, \n",
        "                        name=experiment_name, \n",
        "                        resume=resume_id)\n",
        "\n",
        "        run_name = 'buio/GAN/'+ model_name\n",
        "        artifact_run = run_name +':latest'\n",
        "\n",
        "        artifact = run.use_artifact(artifact_run, type='model')\n",
        "        artifact_dir = artifact.download()\n",
        "        attribute_classifier = tf.keras.models.load_model(os.path.join(artifact_dir, model_name),\n",
        "                                                        custom_objects=custom_objects)\n",
        "\n",
        "        attribute_classifier.compile(loss=keras.losses.BinaryCrossentropy(), metrics='binary_accuracy')\n",
        "        #loss, gen_attr_accuracy = attribute_classifier.evaluate(dataset_df)\n",
        "        #gen_attr_accuracy\n",
        "        run.finish()\n",
        "\n",
        "    else:\n",
        "        #avoid wandb login with wget\n",
        "        attr_cls_path = \"/content/attribute_classifier_model/\"\n",
        "        os.remove(attr_cls_path+\"keras_metadata.pb\")\n",
        "        os.remove(attr_cls_path+\"saved_model.pb\")\n",
        "        os.remove(attr_cls_path+\"variables/variables.data-00000-of-00001\")\n",
        "        os.remove(attr_cls_path+\"variables/variables.index\")\n",
        "        \n",
        "        os.makedirs(attr_cls_path,exist_ok =True)\n",
        "        os.makedirs(attr_cls_path+\"/variables\",exist_ok =True)\n",
        "        \n",
        "\n",
        "        !pip install wget\n",
        "        import wget\n",
        "        wget.download(keras_metadata_url, attr_cls_path+\"keras_metadata.pb\",)\n",
        "        wget.download(saved_model_url, attr_cls_path+\"saved_model.pb\")\n",
        "        wget.download(variables_url, attr_cls_path+\"variables/variables.data-00000-of-00001\")\n",
        "        wget.download(index_url, attr_cls_path+\"variables/variables.index\")\n",
        "\n",
        "        attribute_classifier = tf.keras.models.load_model(attr_cls_path,custom_objects=custom_objects)\n",
        "        attribute_classifier.compile(loss=keras.losses.BinaryCrossentropy(), metrics='binary_accuracy')\n",
        "\n",
        "    attribute_classifier.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPEjfljHsRdY"
      },
      "source": [
        "## Evaluate Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pyohtU1UYaEs"
      },
      "source": [
        "FID_COUNT = 10000\n",
        "train_embeddings, train_attr = generate_dataset_embeddings(dataset_train, FID_COUNT)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rRt6g3URnu_7"
      },
      "source": [
        "def compute_metrics(gan_model):\n",
        "    gen_embeddings, pred_attr = generate_fid_embeddings(gan_model, train_attr, FID_COUNT, use_attributes=True)\n",
        "    fid = compute_fid(train_embeddings, gen_embeddings)\n",
        "    acc, f2 = 0,0\n",
        "    if N_ATTRIBUTES !=0:\n",
        "        acc = np.mean(tf.keras.metrics.binary_accuracy(train_attr[:FID_COUNT], pred_attr[:FID_COUNT]))\n",
        "        f1 = f1_metric_logits(train_attr[:FID_COUNT], pred_attr[:FID_COUNT])\n",
        "    metrics = {'fid':fid, 'f1':f1.numpy(), 'acc':acc}\n",
        "    print(metrics)\n",
        "    return metrics\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hzILYJToHU2U"
      },
      "source": [
        "## Evaluate per epoch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z_KenzIxqlNP"
      },
      "source": [
        "for epoch,gan in zip(EPOCH_RANGE, models):\n",
        "    print(\"Epoch:\",epoch)\n",
        "    compute_metrics(gan)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nPIXJBUPMuTr"
      },
      "source": [
        "# 📱📱 📲Weight Average"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QgmtQskCmju-"
      },
      "source": [
        "gen_weights = []\n",
        "#select epochs to average by slicing\n",
        "for gan in models[1:]:\n",
        "    gen_weights.append(gan.generator.get_weights())\n",
        "\n",
        "avg_weights = list()\n",
        "\n",
        "for weights_list_tuple in zip(*gen_weights):\n",
        "    avg_weights.append(\n",
        "        np.array([np.array(weights).mean(axis=0) \n",
        "        for weights in zip(*weights_list_tuple)]))\n",
        "    \n",
        "\n",
        "avg_gan = gan\n",
        "avg_gan.generator.set_weights(avg_weights)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o1GZbdbom-5l"
      },
      "source": [
        "compute_metrics(avg_gan)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fx0P-87YLsO1"
      },
      "source": [
        "compute_metrics(models[-1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DOKYXZp3Bbnb"
      },
      "source": [
        "# 📸 Pretty plot images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U3jle7WzlIpz"
      },
      "source": [
        "### SAME NOISE DIFFERENT LABEL"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ny0_KXhrOzhB"
      },
      "source": [
        "if N_ATTRIBUTES == 2:\n",
        "    n = tf.random.truncated_normal(shape=(100,))\n",
        "    b = tf.random.truncated_normal(shape=(100,))\n",
        "    noise = np.array(tf.concat([[n for i in range(4)],[b for i in range(4)]], axis = 0))\n",
        "\n",
        "if N_ATTRIBUTES == 10 or N_ATTRIBUTES == 40:\n",
        "    n = tf.random.truncated_normal(shape=(100,))\n",
        "    noise = np.array([n for _ in range(8)])\n",
        "print(noise.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t5X3d9ZwHNAO"
      },
      "source": [
        "if N_ATTRIBUTES == 2:\n",
        "    attr = np.array([[0,0],[0,1],[1,0],[1,1], [0,0],[0,1],[1,0],[1,1]])\n",
        "    \n",
        "if N_ATTRIBUTES == 10:\n",
        "    attr = np.array([[0,0,0,1,1,0,0,0,1,0],                     ## man, smile, nose, cheeck\n",
        "                     [0,0,0,1,0,0,0,0,1,0],                              ## man, nose, cheek\n",
        "                     [1,0,0,1,1,0,1,0,1,0],                              ## man, smile, nose, cheek, wavy\n",
        "                     [1,0,0,1,0,0,1,0,1,0],                              ## man, nose, cheek\n",
        "                     [0,0,0,0,1,0,0,0,1,0],                              ## smile, nose, cheek\n",
        "                     [0,0,0,0,0,0,0,0,1,0],                              ## nose, cheek\n",
        "                     [1,0,0,0,1,0,1,0,1,0],                              ## smile, nose, cheek, wavy\n",
        "                     [1,0,0,0,0,0,1,0,1,0]])                             ## cheek, nose, lipstick, makeup\n",
        "if N_ATTRIBUTES == 40:\n",
        "    attr = np.array([[0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,0,0,0,1,0,1,0,0,0,0,0,1],                     ## man, smile, nose, cheeck\n",
        "                     [0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,0,0,0,0,0,1,0,0,0,0,0,1],                            ## man, nose, cheek\n",
        "                     [0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,1,1,0,0,0,1,0,1,0,0,0,0,0,1],                              ## man, smile, nose, cheek, wavy\n",
        "                     [0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,1,1,0,0,0,0,0,1,0,0,0,0,0,1],                            ## man, nose, cheek\n",
        "                     [0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,1,0,1,0,0,0,0,0,1],                            ## smile, nose, cheek\n",
        "                     [0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,0,1,0,0,0,0,0,1],                            ## man, nose, cheek\n",
        "                     [0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,1,1,0,0,0,1,0,1,0,0,0,0,0,1],                            ## smile, nose, cheek, wavy\n",
        "                     [0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,1,1,0,0,0,0,0,1,0,0,0,0,0,1]])\n",
        "attr.shape\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l-vj_z3Pleut"
      },
      "source": [
        "### DIFFERENT NOISE SAME LABEL"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "syJWxe15lkAv"
      },
      "source": [
        "noise = tf.random.truncated_normal(shape=(8,100))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wCto5chEljsq"
      },
      "source": [
        "if N_ATTRIBUTES == 2:\n",
        "    attr = np.array([[0,1]]*8)\n",
        "\n",
        "if N_ATTRIBUTES == 10:\n",
        "    attr = np.array([[0,0,0,1,1,0,0,0,1,0]]*8)\n",
        "                     \n",
        "if N_ATTRIBUTES == 40:\n",
        "    attr = np.array([[0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,0,0,0,1,0,1,0,0,0,0,0,1]]*8)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3UMmtg9lklM"
      },
      "source": [
        "### PLOT"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WQIl0KbmoyzR"
      },
      "source": [
        "num_img = 8\n",
        "print_rows = 1\n",
        "\n",
        "RANDOM_GEN = False\n",
        "\n",
        "for i in range(1):\n",
        "    if RANDOM_GEN:\n",
        "        random_latent_vectors = tf.random.truncated_normal(shape=(num_img, LATENT_DIM))\n",
        "\n",
        "        if N_ATTRIBUTES !=0:\n",
        "            attr = np.array([attr_list[i] for i in np.random.randint(0,202600, num_img)])\n",
        "            generated_images = gan.generator((random_latent_vectors, attr))\n",
        "        else:\n",
        "            generated_images = gan.generator(random_latent_vectors)\n",
        "    else:\n",
        "        generated_images = gan.generator((noise, attr))\n",
        "\n",
        "    generated_images.numpy()\n",
        "    generated_images = conv_range((-1,1), (0,1))(generated_images)\n",
        "\n",
        "    fig, axes = plt.subplots(print_rows, num_img//print_rows, figsize=(30,20))\n",
        "    #titfont = {'fontname':'Monospace'}\n",
        "    \n",
        "    for m, axis in enumerate(axes):\n",
        "        axis.axis('off')\n",
        "        axis.set_title('\\n'.join([' ☑  '+label*' '.join(LABELS[i].split('_')) for i,label in enumerate(attr[m]) if label]),loc='left',fontsize=15)\n",
        "            \n",
        "        axis.imshow(generated_images[m])\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}