Explaining_in_Style_AttFind.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Explaining_in_Style_AttFind.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "QUZxQ3T_lP-7",
        "cellView": "form"
      },
      "source": [
        "#@title Imports\n",
        "from typing import Optional, Tuple, List\n",
        "import six.moves.cPickle as cPickle\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import requests\n",
        "import tqdm\n",
        "import collections\n",
        "import os\n",
        "import zipfile\n",
        "from matplotlib import pyplot as plt\n",
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "from PIL import ImageFont\n",
        "from io import BytesIO\n",
        "import IPython.display\n"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AdIz_RyzlsEo",
        "cellView": "form"
      },
      "source": [
        "#@title Load the encoder, the discriminator, the generator and the classifier (slow!)\n",
        "path = '/tmp/celab_a_age/' #@param {type: 'string'}\n",
        "bucket_path = 'https://storage.googleapis.com/explaining-in-style/celeb_a_age/'\n",
        "if not os.path.exists(path):\n",
        "  os.mkdir(path)\n",
        "\n",
        "def get_path_from_bucket(bucket_url, path):\n",
        "  r = requests.get(bucket_url)\n",
        "  filename = os.path.split(bucket_url)[-1].replace('.zip', '')\n",
        "  zip_ref = zipfile.ZipFile(BytesIO(r.content))\n",
        "  zip_ref.extractall(path)\n",
        "  return os.path.join(path, filename)\n",
        "\n",
        "num_layers = 14\n",
        "label_size = 2\n",
        "resolution = 256\n",
        "generator = tf.keras.models.load_model(get_path_from_bucket(bucket_path + 'generator.savedmodel.zip', path))\n",
        "encoder = tf.keras.models.load_model(get_path_from_bucket(bucket_path + 'encoder.savedmodel.zip', path))\n",
        "discriminator = tf.keras.models.load_model(get_path_from_bucket(bucket_path + 'discriminator.savedmodel.zip', path))\n",
        "classifier = tf.keras.models.load_model(get_path_from_bucket(bucket_path + 'mobilenet.savedmodel.zip', path))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ldVqDlXSB6SA",
        "cellView": "form"
      },
      "source": [
        "#@title Utils {form-width: '20%'}\n",
        "\n",
        "def make_animation(image: np.ndarray,\n",
        "                   resolution: int,\n",
        "                   figsize: Tuple[int, int] = (20, 8)):\n",
        "  fig = plt.figure(1, figsize=figsize)\n",
        "  _ = plt.gca()\n",
        "\n",
        "  def transpose_image(image):\n",
        "    image_reshape = image.reshape([-1, resolution, resolution, 3])\n",
        "    return image_reshape.transpose([1, 0, 2, 3]).reshape([resolution, -1, 3])\n",
        "  im = plt.imshow(transpose_image(image[:, :resolution, :]),\n",
        "                  interpolation='none')\n",
        "  def animate_func(i):\n",
        "    im.set_array(transpose_image(image[:, resolution*i:resolution*(i+1), :]))\n",
        "    return [im]\n",
        "\n",
        "  animation = matplotlib.animation.FuncAnimation(\n",
        "      fig, animate_func, frames=image.shape[1] // resolution, interval=600)\n",
        "\n",
        "  plt.close(1)\n",
        "  return animation\n",
        "\n",
        "\n",
        "def show_image(image, fmt='png'):\n",
        "  if image.dtype == np.float32:\n",
        "    image = np.uint8(image * 127.5 + 127.5)\n",
        "  if image.shape[0] == 3:\n",
        "    image = np.transpose(image, (1, 2, 0))\n",
        "  bytes_io = BytesIO()\n",
        "  Image.fromarray(image).save(bytes_io, fmt)\n",
        "  IPython.display.display(IPython.display.Image(data=bytes_io.getvalue()))\n",
        "\n",
        "\n",
        "def filter_unstable_images(style_change_effect: np.ndarray,\n",
        "                           effect_threshold: float = 0.3,\n",
        "                           num_indices_threshold: int = 750) -> np.ndarray:\n",
        "  \"\"\"Filters out images which are affected by too many S values.\"\"\"\n",
        "  unstable_images = (\n",
        "      np.sum(np.abs(style_change_effect) > effect_threshold, axis=(1, 2, 3)) >\n",
        "      num_indices_threshold)\n",
        "  style_change_effect[unstable_images] = 0\n",
        "  return style_change_effect\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def call_synthesis(generator: tf.keras.models.Model,\n",
        "                   dlatents_in: tf.Tensor,\n",
        "                   conditioning_in: Optional[tf.Tensor] = None,\n",
        "                   labels_in: Optional[tf.Tensor] = None,\n",
        "                   training: bool = False,\n",
        "                   num_layers: int = 14,\n",
        "                   dlatent_size: int = 512) -> tf.Tensor:\n",
        "  \"\"\"Calls the synthesis.\n",
        "\n",
        "  Args:\n",
        "    dlatents_in: the intermediate latent representation of shape [batch size,\n",
        "      num_layers, dlatent_size].\n",
        "    conditioning_in: Conditioning input to the synthesis network (can be an\n",
        "      image or output from an encoder) of shape [minibatch, channels,\n",
        "      resolution, resolution]. Set to None if unused.\n",
        "    labels_in: of shape [batch_size, label_size]. Set to None if unused.\n",
        "    training: Whether this is a training call.\n",
        "\n",
        "  Returns:\n",
        "    The output images and optional latent vector.\n",
        "\n",
        "  \"\"\"\n",
        "  if labels_in is not None:\n",
        "    zero_labels = tf.zeros_like(labels_in)\n",
        "    dlatents_labels = tf.tile(tf.expand_dims(labels, 1), [1, num_layers, 1])\n",
        "    if dlatent_size > 0:\n",
        "      dlatents_expanded = tf.concat([dlatents_in, dlatents_labels], axis=2)\n",
        "    else:\n",
        "      dlatents_expanded = dlatents_labels\n",
        "  else:\n",
        "    if dlatent_size == 0:\n",
        "      raise ValueError('Dlatents are empty and no labels were provided.')\n",
        "    dlatents_expanded = dlatents_in\n",
        "  # Evaluate synthesis network.\n",
        "  style_vector_blocks, style_vector_torgb = generator.style_vector_calculator(\n",
        "      dlatents_expanded[:, 0], training=training)\n",
        "  if conditioning_in is not None:\n",
        "    network_inputs = (style_vector_blocks, style_vector_torgb,\n",
        "                      conditioning_in)\n",
        "  else:\n",
        "    network_inputs = (style_vector_blocks, style_vector_torgb)\n",
        "  synthesis_results = generator.g_synthesis(network_inputs, training=training)\n",
        "\n",
        "  # Return requested outputs.\n",
        "  return tf.maximum(tf.minimum(synthesis_results, 1), -1)\n",
        "\n",
        "\n",
        "def discriminator_filter(style_change_effect: np.ndarray,\n",
        "                         all_dlatents: np.ndarray,\n",
        "                         generator: tf.keras.models.Model,\n",
        "                         discriminator: tf.keras.models.Model,\n",
        "                         classifier: tf.keras.models.Model,\n",
        "                         sindex: int,\n",
        "                         style_min: float,\n",
        "                         style_max: float,\n",
        "                         class_index: int,\n",
        "                         num_images: int = 10,\n",
        "                         label_size: int = 2,\n",
        "                         change_threshold: float = 0.5,\n",
        "                         shift_size: float = 2,\n",
        "                         effect_threshold: float = 0.2,\n",
        "                         sindex_offset: int = 0) -> bool:\n",
        "  \"\"\"Returns false if changing the style index adds artifacts to the images.\n",
        "\n",
        "  Args:\n",
        "    style_change_effect: A shape of [num_images, 2, style_size, num_classes].\n",
        "      The effect of each change of style on specific direction on each image.\n",
        "    all_dlatents: The dlatents of each image, shape of [num_images,\n",
        "      dlatent_size].\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    discriminator: The discriminator model.\n",
        "    sindex: The style index.\n",
        "    style_min: The style min value in all images.\n",
        "    style_max: The style max value in all images.\n",
        "    class_index: The index of the class to check.\n",
        "    num_images: The number of images to do the disciminator_filter test.\n",
        "    label_size: The label size.\n",
        "    change_threshold: The maximal change allowed in the discriminator\n",
        "      prediction.\n",
        "    shift_size: The size to shift the style index.\n",
        "    effect_threshold: Used for choosing images that the classification was\n",
        "      changed enough.\n",
        "    sindex_offset: The offset of the style index if style_change_effect contains\n",
        "      some of the layers and not all styles.\n",
        "  \"\"\"\n",
        "  for style_sign_index in range(2):\n",
        "    images_idx = ((style_change_effect[:, style_sign_index, sindex,\n",
        "                                       class_index]) >\n",
        "                  effect_threshold).nonzero()[0]\n",
        "\n",
        "    images_idx = images_idx[:num_images]\n",
        "    dlatents = all_dlatents[images_idx]\n",
        "\n",
        "    for i in range(len(images_idx)):\n",
        "      cur_dlatent = dlatents[i:i + 1]\n",
        "      (discriminator_orig, \n",
        "       discriminator_change) = get_discriminator_results_given_dlatent(\n",
        "           dlatent=cur_dlatent,\n",
        "           generator=generator,\n",
        "           discriminator=discriminator,\n",
        "           classifier=classifier,\n",
        "           class_index=class_index,\n",
        "           sindex=sindex + sindex_offset,\n",
        "           s_style_min=style_min,\n",
        "           s_style_max=style_max,\n",
        "           style_direction_index=style_sign_index,\n",
        "           shift_size=shift_size,\n",
        "           label_size=label_size)\n",
        "\n",
        "      if np.abs(discriminator_orig - discriminator_change) > change_threshold:\n",
        "        return False\n",
        "  return True\n",
        "\n",
        "\n",
        "def find_significant_styles_image_fraction(\n",
        "    style_change_effect: np.ndarray,\n",
        "    num_indices: int,\n",
        "    class_index: int,\n",
        "    generator: tf.keras.models.Model,\n",
        "    classifier: tf.keras.models.Model,\n",
        "    all_dlatents: np.ndarray,\n",
        "    style_min: np.ndarray,\n",
        "    style_max: np.ndarray,\n",
        "    effect_threshold: float = 0.2,\n",
        "    min_changed_images_fraction: float = 0.03,\n",
        "    label_size: int = 2,\n",
        "    sindex_offset: int = 0,\n",
        "    discriminator: Optional[tf.keras.models.Model] = None,\n",
        "    discriminator_threshold: float = 0.2) -> List[Tuple[int, int]]:\n",
        "  \"\"\"Returns indices in the style vector which affect the classifier.\n",
        "\n",
        "  Args:\n",
        "    style_change_effect: A shape of [num_images, 2, style_size, num_classes].\n",
        "      The effect of each change of style on specific direction on each image.\n",
        "    num_indices: Number of styles in the result.\n",
        "    class_index: The index of the class to visualize.\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    all_dlatents: The dlatents of each image, shape of [num_images,\n",
        "      dlatent_size].\n",
        "    style_min: An array with the min value for each style index.\n",
        "    style_max: An array with the max value for each style index.\n",
        "    effect_threshold: Minimal change of classifier output to be considered.\n",
        "    min_changed_images_fraction: Minimal fraction of images which are changed.\n",
        "    label_size: The label size.\n",
        "    sindex_offset: The offset of the style index if style_change_effect contains\n",
        "      some of the layers and not all styles.\n",
        "    discriminator: The discriminator model. If None, don't filter style indices.\n",
        "    discriminator_threshold: Used in discriminator_filter to define the maximal\n",
        "      change allowed in the discriminator prediction.\n",
        "    \n",
        "  \"\"\"\n",
        "  effect_positive = np.sum(\n",
        "      style_change_effect[:, :, :, class_index] > effect_threshold, axis=0)\n",
        "  effect_positive = effect_positive.flatten()\n",
        "  all_sindices = []\n",
        "  sindices = np.argsort(effect_positive)[::-1]\n",
        "  if discriminator is not None:\n",
        "    print('Using discriminator...')\n",
        "  for sindex in sindices[:num_indices*2]:\n",
        "    if (effect_positive[sindex] <\n",
        "        min_changed_images_fraction * style_change_effect.shape[0]):\n",
        "      break\n",
        "    if discriminator is not None:\n",
        "      s_index = sindex % style_change_effect.shape[2]\n",
        "      if not discriminator_filter(\n",
        "          style_change_effect,\n",
        "          all_dlatents,\n",
        "          generator,\n",
        "          discriminator,\n",
        "          classifier,\n",
        "          s_index,\n",
        "          style_min[s_index + sindex_offset],\n",
        "          style_max[s_index + sindex_offset],\n",
        "          class_index,\n",
        "          label_size=label_size,\n",
        "          change_threshold=discriminator_threshold,\n",
        "          sindex_offset=sindex_offset):\n",
        "        continue\n",
        "    all_sindices.append(sindex)\n",
        "    if len(all_sindices) == num_indices:\n",
        "      break\n",
        "\n",
        "  return [(x // style_change_effect.shape[2],\n",
        "           (x % style_change_effect.shape[2]) + sindex_offset)\n",
        "          for x in all_sindices]\n",
        "\n",
        "\n",
        "def find_significant_styles(\n",
        "    style_change_effect: np.ndarray,\n",
        "    num_indices: int,\n",
        "    class_index: int,\n",
        "    discriminator: Optional[tf.keras.models.Model],\n",
        "    generator: tf.keras.models.Model,\n",
        "    classifier: tf.keras.models.Model,\n",
        "    all_dlatents: np.ndarray,\n",
        "    style_min: np.ndarray,\n",
        "    style_max: np.ndarray,\n",
        "    max_image_effect: float = 0.2,\n",
        "    label_size: int = 2,\n",
        "    discriminator_threshold: float = 0.2,\n",
        "    sindex_offset: int = 0) -> List[Tuple[int, int]]:\n",
        "  \"\"\"Returns indices in the style vector which affect the classifier.\n",
        "\n",
        "  Args:\n",
        "    style_change_effect: A shape of [num_images, 2, style_size, num_classes].\n",
        "      The effect of each change of style on specific direction on each image.\n",
        "    num_indices: Number of styles in the result.\n",
        "    class_index: The index of the class to visualize.\n",
        "    discriminator: The discriminator model. If None, don't filter style indices.\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    all_dlatents: The dlatents of each image, shape of [num_images,\n",
        "      dlatent_size].\n",
        "    style_min: An array with the min value for each style index.\n",
        "    style_max: An array with the max value for each style index.\n",
        "    max_image_effect: Ignore contributions of styles if the previously found\n",
        "      styles changed the probability of the image by more than this threshold.\n",
        "    label_size: The label size.\n",
        "    discriminator_threshold: Used in discriminator_filter to define the maximal\n",
        "      change allowed in the discriminator prediction.\n",
        "    sindex_offset: The offset of the style index if style_change_effect contains\n",
        "      some of the layers and not all styles.\n",
        "  \"\"\"\n",
        "\n",
        "  num_images = style_change_effect.shape[0]\n",
        "  style_effect_direction = np.maximum(\n",
        "      0, style_change_effect[:, :, :, class_index].reshape((num_images, -1)))\n",
        "\n",
        "  images_effect = np.zeros(num_images)\n",
        "  all_sindices = []\n",
        "  discriminator_removed = []\n",
        "  while len(all_sindices) < num_indices:\n",
        "    next_s = np.argmax(\n",
        "        np.mean(\n",
        "            style_effect_direction[images_effect < max_image_effect], axis=0))\n",
        "    if discriminator is not None:\n",
        "      sindex = next_s % style_change_effect.shape[2]\n",
        "      if sindex == 0:\n",
        "        break\n",
        "      if not discriminator_filter(\n",
        "          style_change_effect=style_change_effect,\n",
        "          all_dlatents=all_dlatents,\n",
        "          generator=generator,\n",
        "          discriminator=discriminator,\n",
        "          classifier=classifier,\n",
        "          sindex=sindex,\n",
        "          style_min=style_min[sindex + sindex_offset],\n",
        "          style_max=style_max[sindex + sindex_offset],\n",
        "          class_index=class_index,\n",
        "          label_size=label_size,\n",
        "          change_threshold=discriminator_threshold,\n",
        "          sindex_offset=sindex_offset):\n",
        "        style_effect_direction[:, next_s] = np.zeros(num_images)\n",
        "        discriminator_removed.append(sindex)\n",
        "        continue\n",
        "\n",
        "    all_sindices.append(next_s)\n",
        "    images_effect += style_effect_direction[:, next_s]\n",
        "    style_effect_direction[:, next_s] = 0\n",
        "\n",
        "  return [(x // style_change_effect.shape[2],\n",
        "           (x % style_change_effect.shape[2]) + sindex_offset)\n",
        "          for x in all_sindices]\n",
        "\n",
        "\n",
        "def _float_features(values):\n",
        "  \"\"\"Returns a float_list from a float / double.\"\"\"\n",
        "  return tf.train.Feature(float_list=tf.train.FloatList(value=values))\n",
        "\n",
        "LAYER_SHAPES = []\n",
        "for dense in generator.style_vector_calculator.style_dense_blocks:\n",
        "  LAYER_SHAPES.append(dense.dense_bias.weights[0].shape[1])\n",
        "\n",
        "\n",
        "def sindex_to_layer_idx_and_index(generator: tf.keras.models.Model,\n",
        "                                  sindex: int) -> Tuple[int, int]:\n",
        "  global LAYER_SHAPES\n",
        "  layer_shapes_cumsum = np.concatenate([[0], np.cumsum(LAYER_SHAPES)])\n",
        "  layer_idx = (layer_shapes_cumsum <= sindex).nonzero()[0][-1]\n",
        "  return layer_idx, sindex - layer_shapes_cumsum[layer_idx]\n",
        "\n",
        "\n",
        "def get_classifier_results(generator: tf.keras.models.Model, \n",
        "                           expanded_dlatent: tf.Tensor,\n",
        "                           use_softmax: bool = False):\n",
        "  image = call_synthesis(generator, expanded_dlatent)\n",
        "  image = tf.transpose(image, (0, 2, 3, 1))\n",
        "  results = classifier(image, training=False)\n",
        "  if use_softmax:\n",
        "    return tf.nn.softmax(results).numpy()[0]\n",
        "  else:\n",
        "    return results.numpy()[0]\n",
        "\n",
        "\n",
        "def draw_on_image(image: np.ndarray, number: float,\n",
        "                  font_file: str,\n",
        "                  font_fill: Tuple[int, int, int] = (0, 0, 255)) -> np.ndarray:\n",
        "  \"\"\"Draws a number on the top left corner of the image.\"\"\"\n",
        "  fnt = ImageFont.truetype(font_file, 20)\n",
        "  out_image = Image.fromarray((image * 127.5 + 127.5).astype(np.uint8))\n",
        "  draw = ImageDraw.Draw(out_image)\n",
        "  draw.multiline_text((10, 10), ('%.3f' % number), font=fnt, fill=font_fill)\n",
        "  return np.array(out_image)\n",
        "\n",
        "\n",
        "def generate_change_image_given_dlatent(\n",
        "    dlatent: np.ndarray,\n",
        "    generator: tf.keras.models.Model,\n",
        "    classifier: Optional[tf.keras.models.Model],\n",
        "    class_index: int,\n",
        "    sindex: int,\n",
        "    s_style_min: float,\n",
        "    s_style_max: float,\n",
        "    style_direction_index: int,\n",
        "    shift_size: float,\n",
        "    label_size: int = 2,\n",
        ") -> Tuple[np.ndarray, float, float]:\n",
        "  \"\"\"Modifies an image given the dlatent on a specific S-index.\n",
        "\n",
        "  Args:\n",
        "    dlatent: The image dlatent, with sape [dlatent_size].\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    classifier: The classifier to visualize.\n",
        "    class_index: The index of the class to visualize.\n",
        "    sindex: The specific style index to visualize.\n",
        "    s_style_min: The minimal value of the style index.\n",
        "    s_style_max: The maximal value of the style index.\n",
        "    style_direction_index: If 0 move s to it's min value otherwise to it's max\n",
        "      value.\n",
        "    shift_size: Factor of the shift of the style vector.\n",
        "    label_size: The size of the label.\n",
        "\n",
        "  Returns:\n",
        "    The image after the style index modification, and the output of\n",
        "    the classifier on this image.\n",
        "  \"\"\"\n",
        "  network_inputs = generator.style_vector_calculator(dlatent)\n",
        "  style_vector = tf.concat(\n",
        "        generator.style_vector_calculator(dlatent, training=False)[0],\n",
        "        axis=1).numpy()\n",
        "  orig_value = style_vector[0, sindex]\n",
        "  target_value = (s_style_min if style_direction_index == 0 else s_style_max)\n",
        "\n",
        "  weight_shift = shift_size * (target_value - orig_value)\n",
        "\n",
        "  layer_idx, in_idx = sindex_to_layer_idx_and_index(generator, sindex)\n",
        "  layer_one_hot = tf.expand_dims(\n",
        "      tf.one_hot(in_idx, network_inputs[0][layer_idx].shape[1]), 0)\n",
        "  network_inputs[0][layer_idx] += (weight_shift * layer_one_hot)\n",
        "  images_out = generator.g_synthesis(network_inputs, training=False)\n",
        "  images_out = tf.maximum(tf.minimum(images_out, 1), -1)\n",
        "  change_image = tf.transpose(images_out, [0, 2, 3, 1])\n",
        "  result = classifier(change_image, training=False)\n",
        "  change_prob = tf.nn.softmax(result).numpy()[0, class_index]\n",
        "  return change_image, change_prob\n",
        "\n",
        "\n",
        "\n",
        "def get_discriminator_results_given_dlatent(\n",
        "    dlatent: np.ndarray,\n",
        "    generator: tf.keras.models.Model,\n",
        "    discriminator: tf.keras.models.Model,\n",
        "    classifier: tf.keras.models.Model,\n",
        "    class_index: int,\n",
        "    sindex: int,\n",
        "    s_style_min: float,\n",
        "    s_style_max: float,\n",
        "    style_direction_index: int,\n",
        "    shift_size: float = 2,\n",
        "    label_size: int = 2,\n",
        ") -> Tuple[float, float]:\n",
        "  \"\"\"Modifies an image given the dlatent on a specific S-index.\n",
        "\n",
        "  Args:\n",
        "    dlatent: The image dlatent, with sape [dlatent_size].\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    class_index: The index of the class to visualize.\n",
        "    sindex: The specific style index to visualize.\n",
        "    s_style_min: The minimal value of the style index.\n",
        "    s_style_max: The maximal value of the style index.\n",
        "    style_direction_index: If 0 move s to it's min value otherwise to it's max\n",
        "      value.\n",
        "    shift_size: Factor of the shift of the style vector.\n",
        "    label_size: The size of the label.\n",
        "\n",
        "  Returns:\n",
        "    The discriminator before and after.\n",
        "  \"\"\"\n",
        "  network_inputs = generator.style_vector_calculator(dlatent)\n",
        "  images_out = generator.g_synthesis(network_inputs, training=False)\n",
        "  images_out = tf.maximum(tf.minimum(images_out, 1), -1)\n",
        "  labels = tf.constant(dlatent[:, -label_size:], dtype=tf.float32)\n",
        "  discriminator_before = discriminator([images_out, labels], training=False)\n",
        "  # I am not using the classifier output here, because it is only one.\n",
        "  change_image, _ = (\n",
        "      generate_change_image_given_dlatent(dlatent, generator, classifier,\n",
        "                                          class_index, sindex,\n",
        "                                          s_style_min, s_style_max,\n",
        "                                          style_direction_index, shift_size,\n",
        "                                          label_size))\n",
        "  labels = tf.nn.softmax(classifier(change_image, training=False))\n",
        "  change_image_for_disc = tf.transpose(change_image, (0, 3, 1, 2))\n",
        "  discriminator_after = discriminator([change_image_for_disc, labels], \n",
        "                                      training=False)\n",
        "  return (discriminator_before, discriminator_after)\n",
        "\n",
        "\n",
        "def generate_images_given_dlatent(\n",
        "    dlatent: np.ndarray,\n",
        "    generator: tf.keras.models.Model,\n",
        "    classifier: Optional[tf.keras.models.Model],\n",
        "    class_index: int,\n",
        "    sindex: int,\n",
        "    s_style_min: float,\n",
        "    s_style_max: float,\n",
        "    style_direction_index: int,\n",
        "    font_file: Optional[str],\n",
        "    shift_size: float = 2,\n",
        "    label_size: int = 2,\n",
        "    draw_results_on_image: bool = True,\n",
        "    resolution: int = 256,\n",
        ") -> Tuple[np.ndarray, float, float, float, float]:\n",
        "  \"\"\"Modifies an image given the dlatent on a specific S-index.\n",
        "\n",
        "  Args:\n",
        "    dlatent: The image dlatent, with sape [dlatent_size].\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    classifier: The classifier to visualize.\n",
        "    class_index: The index of the class to visualize.\n",
        "    sindex: The specific style index to visualize.\n",
        "    s_style_min: The minimal value of the style index.\n",
        "    s_style_max: The maximal value of the style index.\n",
        "    style_direction_index: If 0 move s to it's min value otherwise to it's max\n",
        "      value.\n",
        "    font_file: A path to the font file for writing the probability on the image.\n",
        "    shift_size: Factor of the shift of the style vector.\n",
        "    label_size: The size of the label.\n",
        "    draw_results_on_image: Whether to draw the classifier outputs on the images.\n",
        "\n",
        "  Returns:\n",
        "    The image before and after the style index modification, and the outputs of\n",
        "    the classifier before and after the\n",
        "    modification.\n",
        "  \"\"\"\n",
        "  network_inputs = generator.style_vector_calculator(dlatent)\n",
        "  result_image = np.zeros((resolution, 2 * resolution, 3), np.uint8)\n",
        "  images_out = generator.g_synthesis(network_inputs, training=False)\n",
        "  images_out = tf.maximum(tf.minimum(images_out, 1), -1)\n",
        "  base_image = tf.transpose(images_out, [0, 2, 3, 1])\n",
        "  result = classifier(base_image, training=False)\n",
        "  base_prob = tf.nn.softmax(result).numpy()[0, class_index]\n",
        "  if draw_results_on_image:\n",
        "    result_image[:, :resolution, :] = draw_on_image(\n",
        "        base_image[0].numpy(), base_prob, font_file)\n",
        "  else:\n",
        "    result_image[:, :resolution, :] = (base_image[0].numpy() * 127.5 +\n",
        "                                       127.5).astype(np.uint8)\n",
        "\n",
        "  change_image, change_prob = (\n",
        "      generate_change_image_given_dlatent(dlatent, generator, classifier,\n",
        "                                          class_index, sindex,\n",
        "                                          s_style_min, s_style_max,\n",
        "                                          style_direction_index, shift_size,\n",
        "                                          label_size))\n",
        "  if draw_results_on_image:\n",
        "    result_image[:, resolution:, :] = draw_on_image(\n",
        "        change_image[0].numpy(), change_prob, font_file)\n",
        "  else:\n",
        "    result_image[:, resolution:, :] = (\n",
        "        np.maxiumum(np.minimum(change_image[0].numpy(), 1), -1) * 127.5 +\n",
        "                                               127.5).astype(np.uint8)\n",
        "\n",
        "  return (result_image, change_prob, base_prob)\n",
        "\n",
        "\n",
        "\n",
        "def visualize_style(generator: tf.keras.models.Model,\n",
        "                    classifier: tf.keras.models.Model,\n",
        "                    all_dlatents: np.ndarray,\n",
        "                    style_change_effect: np.ndarray,\n",
        "                    style_min: np.ndarray,\n",
        "                    style_max: np.ndarray,\n",
        "                    sindex: int,\n",
        "                    style_direction_index: int,\n",
        "                    max_images: int,\n",
        "                    shift_size: float,\n",
        "                    font_file: str,\n",
        "                    label_size: int = 2,\n",
        "                    class_index: int = 0,\n",
        "                    effect_threshold: float = 0.3,\n",
        "                    seed: Optional[int] = None,\n",
        "                    allow_both_directions_change: bool = False,\n",
        "                    draw_results_on_image: bool = True) -> np.ndarray:\n",
        "  \"\"\"Returns an image visualizing the effect of a specific S-index.\n",
        "\n",
        "  Args:\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    classifier: The classifier to visualize.\n",
        "    all_dlatents: An array with shape [num_images, dlatent_size].\n",
        "    style_change_effect: A shape of [num_images, 2, style_size, num_classes].\n",
        "      The effect of each change of style on specific direction on each image.\n",
        "    style_min: The minimal value of each style, with shape [style_size].\n",
        "    style_max: The maximal value of each style, with shape [style_size].\n",
        "    sindex: The specific style index to visualize.\n",
        "    style_direction_index: If 0 move s to its min value otherwise to its max\n",
        "      value.\n",
        "    max_images: Maximal number of images to visualize.\n",
        "    shift_size: Factor of the shift of the style vector.\n",
        "    font_file: A path to the font file for writing the probability on the image.\n",
        "    label_size: The size of the label.\n",
        "    class_index: The index of the class to visualize.\n",
        "    effect_threshold: Choose images whose effect was at least this number.\n",
        "    seed: If given, use this as a seed to the random shuffling of the images.\n",
        "    allow_both_directions_change: Whether to allow both increasing and\n",
        "      decreasing the classifiaction (used for age).\n",
        "    draw_results_on_image: Whether to draw the classifier outputs on the images.\n",
        "  \"\"\"\n",
        "\n",
        "  # Choose the dlatent indices to visualize\n",
        "  if allow_both_directions_change:\n",
        "    images_idx = (np.abs(style_change_effect[:, style_direction_index, sindex,\n",
        "                                             class_index]) >\n",
        "                  effect_threshold).nonzero()[0]\n",
        "  else:\n",
        "    images_idx = ((style_change_effect[:, style_direction_index, sindex,\n",
        "                                       class_index]) >\n",
        "                  effect_threshold).nonzero()[0]\n",
        "  if images_idx.size == 0:\n",
        "    return np.array([])\n",
        "\n",
        "  if seed is not None:\n",
        "    np.random.seed(seed)\n",
        "  np.random.shuffle(images_idx)\n",
        "  images_idx = images_idx[:min(max_images*10, len(images_idx))]\n",
        "  dlatents = all_dlatents[images_idx]\n",
        "\n",
        "  result_images = []\n",
        "  for i in range(len(images_idx)):\n",
        "    cur_dlatent = dlatents[i:i + 1]\n",
        "    (result_image, base_prob, change_prob) = generate_images_given_dlatent(\n",
        "         dlatent=cur_dlatent,\n",
        "         generator=generator,\n",
        "         classifier=classifier,\n",
        "         class_index=class_index,\n",
        "         sindex=sindex,\n",
        "         s_style_min=style_min[sindex],\n",
        "         s_style_max=style_max[sindex],\n",
        "         style_direction_index=style_direction_index,\n",
        "         font_file=font_file,\n",
        "         shift_size=shift_size,\n",
        "         label_size=label_size,\n",
        "         draw_results_on_image=draw_results_on_image)\n",
        "\n",
        "    if np.abs(change_prob - base_prob) < effect_threshold:\n",
        "      continue\n",
        "    result_images.append(result_image)\n",
        "    if len(result_images) == max_images:\n",
        "      break\n",
        "\n",
        "  if len(result_images) < 3:\n",
        "    # No point in returning results with very little images\n",
        "    return np.array([])\n",
        "  return np.concatenate(result_images[:max_images], axis=0)\n",
        "\n",
        "\n",
        "def visualize_style_by_distance_in_s(\n",
        "    generator: tf.keras.models.Model,\n",
        "    classifier: tf.keras.models.Model,\n",
        "    all_dlatents: np.ndarray,\n",
        "    all_style_vectors_distances: np.ndarray,\n",
        "    style_min: np.ndarray,\n",
        "    style_max: np.ndarray,\n",
        "    sindex: int,\n",
        "    style_sign_index: int,\n",
        "    max_images: int,\n",
        "    shift_size: float,\n",
        "    font_file: str,\n",
        "    label_size: int = 2,\n",
        "    class_index: int = 0,\n",
        "    draw_results_on_image: bool = True,\n",
        "    effect_threshold: float = 0.1) -> np.ndarray:\n",
        "  \"\"\"Returns an image visualizing the effect of a specific S-index.\n",
        "\n",
        "  Args:\n",
        "    generator: The generator model. Either StyleGAN or GLO.\n",
        "    classifier: The classifier to visualize.\n",
        "    all_dlatents: An array with shape [num_images, dlatent_size].\n",
        "    all_style_vectors_distances: A shape of [num_images, style_size, 2].\n",
        "      The distance each style from the min and max values on each image.\n",
        "    style_min: The minimal value of each style, with shape [style_size].\n",
        "    style_max: The maximal value of each style, with shape [style_size].\n",
        "    sindex: The specific style index to visualize.\n",
        "    style_sign_index: If 0 move s to its min value otherwise to its max\n",
        "      value.\n",
        "    max_images: Maximal number of images to visualize.\n",
        "    shift_size: Factor of the shift of the style vector.\n",
        "    font_file: A path to the font file for writing the probability on the image.\n",
        "    label_size: The size of the label.\n",
        "    class_index: The index of the class to visualize.\n",
        "    draw_results_on_image: Whether to draw the classifier outputs on the images.\n",
        "  \"\"\"\n",
        "\n",
        "  # Choose the dlatent indices to visualize\n",
        "  images_idx = np.argsort(\n",
        "      all_style_vectors_distances[:, sindex, style_sign_index])[::-1]\n",
        "  if images_idx.size == 0:\n",
        "    return np.array([])\n",
        "\n",
        "  images_idx = images_idx[:min(max_images*10, len(images_idx))]\n",
        "  dlatents = all_dlatents[images_idx]\n",
        "\n",
        "  result_images = []\n",
        "  for i in range(len(images_idx)):\n",
        "    cur_dlatent = dlatents[i:i + 1]\n",
        "    (result_image, change_prob, base_prob) = generate_images_given_dlatent(\n",
        "         dlatent=cur_dlatent,\n",
        "         generator=generator,\n",
        "         classifier=classifier,\n",
        "         class_index=class_index,\n",
        "         sindex=sindex,\n",
        "         s_style_min=style_min[sindex],\n",
        "         s_style_max=style_max[sindex],\n",
        "         style_direction_index=style_sign_index,\n",
        "         font_file=font_file,\n",
        "         shift_size=shift_size,\n",
        "         label_size=label_size,\n",
        "         draw_results_on_image=draw_results_on_image)\n",
        "    if (change_prob - base_prob) < effect_threshold:\n",
        "      continue\n",
        "    result_images.append(result_image)\n",
        "\n",
        "\n",
        "  if len(result_images) < 3:\n",
        "    # No point in returning results with very little images\n",
        "    return np.array([])\n",
        "  return np.concatenate(result_images[:max_images], axis=0)\n"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rO5zl3a6mJDB"
      },
      "source": [
        "# Latents extraction: \n",
        "We don't provide the dataset_provider, so we will load the dlatents from a precomputed np.array.\n",
        "Here we present how to run the encoder on one image, to calculate one dlatent (we pre-calculate for 250)\n",
        "\n",
        "`TODO:` remove this code (or replace the image with a better reconstruction)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dNR2Xn5HmH9Y",
        "cellView": "form"
      },
      "source": [
        "#@title Load the precomputed dlatents (already concatenated to the labels)\n",
        "dlatents_url = bucket_path + 'dlatents.pkl'\n",
        "r = requests.get(dlatents_url)\n",
        "dlatents = cPickle.loads(r.content)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xvHqYl9THpC",
        "cellView": "form"
      },
      "source": [
        "#@title Generate image from the dlatent\n",
        "\n",
        "dlatent_index =  4#@param {type: \"integer\"}\n",
        "\n",
        "expanded_dlatent_tmp = tf.tile(\n",
        "      tf.expand_dims(dlatents[dlatent_index][1], 1),\n",
        "      [1, num_layers, 1])\n",
        "rec_image = call_synthesis(generator,\n",
        "    expanded_dlatent_tmp,\n",
        "    num_layers=num_layers)\n",
        "image = rec_image.numpy()[0]\n",
        "show_image(image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D-QDwtgbpWaE"
      },
      "source": [
        "## Run the Generation step of AttFind\n",
        "Do not run if loaded from precomputed dlatents."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o8zCmWKwpaN0",
        "cellView": "form"
      },
      "source": [
        "#@title Generate SSpace per index (min and max) \n",
        "values_per_index = collections.defaultdict(list)\n",
        "for _, dlatent in dlatents:\n",
        "  # Get the style vector: \n",
        "  s_img = tf.concat(generator.style_vector_calculator(\n",
        "      dlatent, training=False)[0], axis=1).numpy()[0]\n",
        "  for i, s_val in enumerate(s_img):\n",
        "    values_per_index[i].append(s_val)\n",
        "\n",
        "values_per_index = dict(values_per_index)\n",
        "s_indices_num = len(values_per_index.keys())\n",
        "minimums = [min(values_per_index[i]) for i in range(s_indices_num)] \n",
        "maximums = [max(values_per_index[i]) for i in range(s_indices_num)] "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fXBJJJBSpu37",
        "cellView": "form"
      },
      "source": [
        "#@title SSpace calculation (this is also very heavy, skip to next cell to load precomputed ones) {form-width: '20%'}\n",
        "\n",
        "s_shift_size = 1 # @param\n",
        "data_path = '/tmp/celab_a_age/examples_1.tfrecord' #@param {type: 'string'}\n",
        "\n",
        "with tf.io.TFRecordWriter(data_path) as writer:\n",
        "  for dlatent_index, dlatent in dlatents: \n",
        "    print(dlatent_index)\n",
        "    expanded_dlatent = tf.tile(\n",
        "          tf.expand_dims(dlatent, 1),\n",
        "          [1, num_layers, 1])\n",
        "    base_prob = get_classifier_results(generator, expanded_dlatent)\n",
        "    classifier_results = []\n",
        "    for sindex in tqdm.tqdm(range(0, s_indices_num)):\n",
        "      layer_idx, weight_idx = sindex_to_layer_idx_and_index(generator, sindex)\n",
        "      layer = generator.style_vector_calculator.style_dense_blocks[layer_idx]\n",
        "      layer_size = layer.dense_bias.weights[0].shape[1]\n",
        "      # Get the style vector.\n",
        "      s_vals = tf.concat(\n",
        "          generator.style_vector_calculator(dlatent, training=False)[0], \n",
        "          axis=1).numpy()[0]\n",
        "      s_shift_down = (minimums[sindex] - s_vals[sindex]) * s_shift_size\n",
        "      s_shift_up = (maximums[sindex] - s_vals[sindex]) * s_shift_size\n",
        "      s_shift_d = s_shift_down * tf.expand_dims(tf.one_hot(weight_idx, \n",
        "                                                          layer_size), axis=0)\n",
        "      layer.dense_bias.weights[0].assign_add(s_shift_d)\n",
        "      classifier_results.extend(\n",
        "          get_classifier_results(generator, expanded_dlatent) - base_prob)\n",
        "      layer.dense_bias.weights[0].assign_add(-s_shift_d)\n",
        "      s_shift_u = s_shift_up * tf.expand_dims(\n",
        "          tf.one_hot(weight_idx, layer_size), axis=0)\n",
        "      layer.dense_bias.weights[0].assign_add(s_shift_u)\n",
        "      classifier_results.extend(\n",
        "          get_classifier_results(generator, expanded_dlatent) - base_prob)\n",
        "      layer.dense_bias.weights[0].assign_add(-s_shift_u)\n",
        "    \n",
        "    feature = {}\n",
        "    feature['base_prob'] = _float_features(base_prob.flatten())\n",
        "    feature['dlatent'] = _float_features(dlatent.flatten())\n",
        "    feature['result'] = _float_features(np.array(classifier_results).flatten())\n",
        "    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n",
        "    writer.write(example_proto.SerializeToString())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s_08W22c1ODC"
      },
      "source": [
        "## Run the extraction step of AttFind"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RitFLovd0BBf",
        "cellView": "form"
      },
      "source": [
        "#@title Load effect data from the tfrecord {form-width: '20%'}\n",
        "examples_url = bucket_path + 'examples_1.tfrecord'\n",
        "r = requests.get(examples_url)\n",
        "data_path = path + 'examples_1.tfrecord'\n",
        "open(data_path, 'wb').write(r.content)\n",
        "num_classes = 2\n",
        "print(f'Loaded dataset: {data_path}')\n",
        "table = tf.data.TFRecordDataset([data_path])\n",
        "# Read sspace tfrecord unwrapped:\n",
        "style_change_effect = []\n",
        "dlatents = []\n",
        "base_probs = []\n",
        "for raw_record in table:\n",
        "  example = tf.train.Example()\n",
        "  example.ParseFromString(raw_record.numpy())\n",
        "  dlatents.append(\n",
        "      np.array(example.features.feature['dlatent'].float_list.value))\n",
        "  seffect = np.array(\n",
        "      example.features.feature['result'].float_list.value).reshape(\n",
        "          (-1, 2, num_classes))\n",
        "  style_change_effect.append(seffect.transpose([1, 0, 2]))\n",
        "  base_probs.append(\n",
        "      np.array(example.features.feature['base_prob'].float_list.value))\n",
        "\n",
        "style_change_effect = np.array(style_change_effect)\n",
        "dlatents = np.array(dlatents)\n",
        "W_values, style_change_effect, base_probs = dlatents, style_change_effect, np.array(base_probs)\n",
        "\n",
        "\n",
        "style_change_effect = filter_unstable_images(style_change_effect, effect_threshold=2)\n",
        "\n",
        "all_style_vectors = tf.concat(generator.style_vector_calculator(W_values, training=False)[0], axis=1).numpy()\n",
        "style_min = np.min(all_style_vectors, axis=0)\n",
        "style_max = np.max(all_style_vectors, axis=0)\n",
        "\n",
        "all_style_vectors_distances = np.zeros((all_style_vectors.shape[0], all_style_vectors.shape[1], 2))\n",
        "all_style_vectors_distances[:,:, 0] = all_style_vectors - np.tile(style_min, (all_style_vectors.shape[0], 1))\n",
        "all_style_vectors_distances[:,:, 1] = np.tile(style_max, (all_style_vectors.shape[0], 1)) - all_style_vectors\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UsoPIHtHsAeS",
        "cellView": "form"
      },
      "source": [
        "#@title Split by class\n",
        "all_labels = np.argmax(base_probs, axis=1)\n",
        "style_effect_classes = {}\n",
        "W_classes = {}\n",
        "style_vectors_distances_classes = {}\n",
        "all_style_vectors_classes = {}\n",
        "for img_ind in range(label_size):\n",
        "  img_inx = np.array([i for i in range(all_labels.shape[0]) \n",
        "  if all_labels[i] == img_ind])\n",
        "  curr_style_effect = np.zeros((len(img_inx), style_change_effect.shape[1], \n",
        "                                style_change_effect.shape[2], style_change_effect.shape[3]))\n",
        "  curr_w = np.zeros((len(img_inx), W_values.shape[1]))\n",
        "  curr_style_vector_distances = np.zeros((len(img_inx), style_change_effect.shape[2], 2))\n",
        "  for k, i in enumerate(img_inx):\n",
        "    curr_style_effect[k, :, :] = style_change_effect[i, :, :, :]\n",
        "    curr_w[k, :] = W_values[i, :]\n",
        "    curr_style_vector_distances[k, :, :] = all_style_vectors_distances[i, :, :]\n",
        "  style_effect_classes[img_ind] = curr_style_effect\n",
        "  W_classes[img_ind] = curr_w\n",
        "  style_vectors_distances_classes[img_ind] = curr_style_vector_distances\n",
        "  all_style_vectors_classes[img_ind] = all_style_vectors[img_inx]\n",
        "  print(f'Class {img_ind}, {len(img_inx)} images.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9KRBsoN7zfFf",
        "cellView": "form"
      },
      "source": [
        "#@title Significant S values - combined {form-width: '20%'}\n",
        "label_size_clasifier = 2 #@param\n",
        "num_indices =  8 #@param\n",
        "effect_threshold = 0.2 #@param\n",
        "use_discriminator = False #@param {type: 'boolean'}\n",
        "discriminator_model = discriminator if use_discriminator else None\n",
        "s_indices_and_signs_dict = {}\n",
        "\n",
        "for class_index in [0, 1]:\n",
        "  split_ind = 1 - class_index\n",
        "  all_s = style_effect_classes[split_ind]\n",
        "  all_w = W_classes[split_ind]\n",
        "\n",
        "  # Find s indicies\n",
        "  s_indices_and_signs = find_significant_styles(\n",
        "    style_change_effect=all_s,\n",
        "    num_indices=num_indices,\n",
        "    class_index=class_index,\n",
        "    discriminator=discriminator_model,\n",
        "    generator=generator,\n",
        "    classifier=classifier,\n",
        "    all_dlatents=all_w,\n",
        "    style_min=style_min,\n",
        "    style_max=style_max,\n",
        "    max_image_effect=effect_threshold*5,\n",
        "    label_size=label_size_clasifier,\n",
        "    discriminator_threshold=0.2,\n",
        "    sindex_offset=0)\n",
        "\n",
        "  s_indices_and_signs_dict[class_index] = s_indices_and_signs\n",
        "\n",
        "# Combine the style indicies for the two classes.\n",
        "sindex_class_0 = [sindex for _, sindex in s_indices_and_signs_dict[0]]\n",
        "\n",
        "all_sindex_joined_class_0 = [(1 - direction, sindex) for direction, sindex in \n",
        "                             s_indices_and_signs_dict[1] if sindex not in sindex_class_0]\n",
        "all_sindex_joined_class_0 += s_indices_and_signs_dict[0]\n",
        "\n",
        "scores = []\n",
        "for direction, sindex in all_sindex_joined_class_0:\n",
        "  other_direction = 1 if direction == 0 else 0\n",
        "  curr_score = np.mean(style_change_effect[:, direction, sindex, 0]) + np.mean(style_change_effect[:, other_direction, sindex, 1])\n",
        "  scores.append(curr_score)\n",
        "\n",
        "s_indices_and_signs = [all_sindex_joined_class_0[i] for i in np.argsort(scores)[::-1]]\n",
        "\n",
        "print('Directions and style indices for moving from class 1 to class 0 = ', s_indices_and_signs[:num_indices])\n",
        "print('Use the other direction to move for class 0 to 1.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JCIUX9504VMV",
        "cellView": "form"
      },
      "source": [
        "#@title Significant S values - multi class one vs. all (this can also be used as one vs. one, using split_by_class=True and define class_index_from) {form-width: '20%'}\n",
        "label_size_clasifier = 2 #@param {type:\"integer\"}\n",
        "num_indices =  8#@param\n",
        "effect_threshold = 0.1 #@param\n",
        "use_discriminator = False #@param {type: 'boolean'}\n",
        "class_index_to = 0 #@param {type:\"integer\"}\n",
        "discriminator_model = discriminator if use_discriminator else None\n",
        "split_by_class = False #@param {type: 'boolean'}\n",
        "class_index_from = 1 #@param {type:\"integer\"}\n",
        "\n",
        "\n",
        "if split_by_class:\n",
        "  split_ind = class_index_from\n",
        "  all_s = style_effect_classes[split_ind]\n",
        "  all_w = W_classes[split_ind]\n",
        "else:\n",
        "  all_s = style_change_effect\n",
        "  all_w = W_values\n",
        "\n",
        "\n",
        "# Find s indicies\n",
        "s_indices_and_signs = find_significant_styles(\n",
        "  style_change_effect=all_s,\n",
        "  num_indices=num_indices,\n",
        "  class_index=class_index_to,\n",
        "  discriminator=discriminator_model,\n",
        "  generator=generator,\n",
        "  classifier=classifier,\n",
        "  all_dlatents=all_w,\n",
        "  style_min=style_min,\n",
        "  style_max=style_max,\n",
        "  max_image_effect=effect_threshold*5,\n",
        "  label_size=label_size_clasifier,\n",
        "  discriminator_threshold=0.2,\n",
        "  sindex_offset=0)\n",
        "  \n",
        "print(f'Directions and style indices for moving from all classes to class {class_index_to} = ', s_indices_and_signs[:num_indices])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fB7vHm5R5YW7",
        "cellView": "form"
      },
      "source": [
        "#@title Visualize s-index {form-width: '20%'}\n",
        "\n",
        "max_images = 10 #@param\n",
        "sindex =   5300#@param\n",
        "class_index = 0#@param {type: \"integer\"} \n",
        "shift_sign = \"1\" #@param [0, 1]\n",
        "wsign_index = int(shift_sign)\n",
        "\n",
        "shift_size =  1#@param\n",
        "effect_threshold =  0.2#@param\n",
        "split_by_class = True #@param {type:\"boolean\"}\n",
        "select_images_by_s_distance = True #@param {type:\"boolean\"}\n",
        "draw_results_on_image = True #@param {type:\"boolean\"}\n",
        "\n",
        "if split_by_class:\n",
        "  split_ind = 1 if class_index == 0 else 0\n",
        "  all_s = style_effect_classes[split_ind]\n",
        "  all_w = W_classes[split_ind]\n",
        "  all_s_distances = style_vectors_distances_classes[split_ind]\n",
        "else:\n",
        "  all_s = style_change_effect\n",
        "  all_w = W_values\n",
        "  all_s_distances = all_style_vectors_distances\n",
        "\n",
        "font_file = '/tmp/arialuni.ttf'\n",
        "if not os.path.exists(font_file):\n",
        "  r = requests.get('https://storage.googleapis.com/google-code-archive-downloads/v2/code.google.com/ipwn/arialuni.ttf')\n",
        "  open(font_file, 'wb').write(r.content)\n",
        "\n",
        "if not select_images_by_s_distance:\n",
        "  yy = visualize_style(generator, \n",
        "                       classifier,\n",
        "                       all_w,\n",
        "                       all_s,\n",
        "                       style_min,\n",
        "                       style_max,\n",
        "                       sindex,\n",
        "                       wsign_index,\n",
        "                       max_images=max_images,\n",
        "                       shift_size=shift_size,\n",
        "                       font_file=font_file,\n",
        "                       label_size=label_size,\n",
        "                       class_index=class_index,\n",
        "                       effect_threshold=effect_threshold,\n",
        "                       draw_results_on_image=draw_results_on_image)\n",
        "    \n",
        "else:\n",
        "  yy = visualize_style_by_distance_in_s(\n",
        "    generator,\n",
        "    classifier,\n",
        "    all_w,\n",
        "    all_s_distances,\n",
        "    style_min,\n",
        "    style_max,\n",
        "    sindex,\n",
        "    wsign_index,\n",
        "    max_images=max_images,\n",
        "    shift_size=shift_size,\n",
        "    font_file=font_file,\n",
        "    label_size=label_size,\n",
        "    class_index=class_index,\n",
        "    effect_threshold=effect_threshold,\n",
        "    draw_results_on_image=draw_results_on_image)\n",
        "\n",
        "if yy.size > 0:\n",
        "  show_image(yy)\n",
        "else:\n",
        "  print('no images found')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bM4glIz6lqcH",
        "cellView": "form"
      },
      "source": [
        "#@title Show animation {form-width: '20%'}\n",
        "\n",
        "import matplotlib.animation\n",
        "from IPython.display import HTML\n",
        "\n",
        "ani = make_animation(yy, resolution)\n",
        "\n",
        "HTML(ani.to_jshtml())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "jJFeg632yhCM"
      },
      "source": [
        "#@title Show the 4 top attributes - as displayed in Fig.4 (b)\n",
        "\n",
        "draw_probabilities_on_image = True #@param {type: \"boolean\"}\n",
        "index_to_naming = {0: \"Skin Pigminatation\", 1: \"Eyebrow Thickness\", 2: \"Add/Remove Glasses\", 3: \"Dark/White Hair\"}\n",
        "images_list = [[0, 4], [16, 17], [18, 18], [3, 14]]\n",
        "max_images = 20\n",
        "shift_sizes = [(2, 1.5),(1, 1),(1, 1),(1.5, 2)]\n",
        "effect_threshold = 0.05\n",
        "font_file = '/tmp/arialuni.ttf'\n",
        "if not os.path.exists(font_file):\n",
        "  gfile.Copy('/google_src/head/depot/google3/googledata/third_party/fonts/ascender/arialuni.ttf', font_file)\n",
        "\n",
        "for i, (direction, sindex) in enumerate(s_indices_and_signs[:4]):\n",
        "  images_s = np.zeros((resolution * 2, resolution * 2, 3)).astype(np.uint8)\n",
        "  for d in [direction, 1 - direction]:\n",
        "    # Take only images from the offsite class\n",
        "    class_index = 0 if d == direction else 1\n",
        "    split_ind = 1 if d == direction else 0\n",
        "    all_s = style_effect_classes[split_ind]\n",
        "    all_w = W_classes[split_ind]\n",
        "    all_s_distances = style_vectors_distances_classes[split_ind]\n",
        "    # Generate images\n",
        "    yy = visualize_style_by_distance_in_s(\n",
        "      generator,\n",
        "      classifier,\n",
        "      all_w,\n",
        "      all_s_distances,\n",
        "      style_min,\n",
        "      style_max,\n",
        "      sindex,\n",
        "      d,\n",
        "      max_images=max_images,\n",
        "      shift_size=shift_sizes[i][class_index],\n",
        "      font_file=font_file,\n",
        "      label_size=label_size,\n",
        "      class_index=class_index,\n",
        "      effect_threshold=effect_threshold,\n",
        "      draw_results_on_image=draw_probabilities_on_image)\n",
        "    for n in range(2):\n",
        "      images_s[n * resolution: (n + 1) * resolution, class_index * resolution: (class_index + 1) * resolution, :] = yy[(images_list[i][class_index]) * resolution: (images_list[i][class_index] + 1) * resolution, n * resolution: (n + 1) * resolution, :]\n",
        "  print(f'Attribute {i} {index_to_naming[i]}: \\n(Original images are on the first row, the probabilities displayed are for the other class - left column for being old, write column for being young)')\n",
        "  show_image(images_s)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}