StyleSDF_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleSDF_demo.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/royorel/StyleSDF/blob/main/StyleSDF_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#StyleSDF Demo\n",
        "\n",
        "This Colab notebook demonstrates the capabilities of the StyleSDF 3D-aware GAN architecture proposed in our paper.\n",
        "\n",
        "This colab generates images with their correspondinig 3D meshes"
      ],
      "metadata": {
        "id": "b86fxLSo1gqI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, let's download the github repository and install all dependencies."
      ],
      "metadata": {
        "id": "8wSAkifH2Bk5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/royorel/StyleSDF.git\n",
        "%cd StyleSDF\n",
        "!pip3 install -r requirements.txt"
      ],
      "metadata": {
        "id": "1GTFRig12CuH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And install pytorch3D..."
      ],
      "metadata": {
        "id": "GwFJ8oLY4i8d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U fvcore\n",
        "import sys\n",
        "import torch\n",
        "pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
        "version_str=\"\".join([\n",
        "    f\"py3{sys.version_info.minor}_cu\",\n",
        "    torch.version.cuda.replace(\".\",\"\"),\n",
        "    f\"_pyt{pyt_version_str}\"\n",
        "])\n",
        "!pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html"
      ],
      "metadata": {
        "id": "zl3Vpddz3ols"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's download the pretrained models for FFHQ and AFHQ."
      ],
      "metadata": {
        "id": "OgMcslVS5vbC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python download_models.py"
      ],
      "metadata": {
        "id": "r1iDkz7r5wnO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we import libraries and set options.\n",
        "\n",
        "Note: this might take a while (approx. 1-2 minutes) since CUDA kernels need to be compiled."
      ],
      "metadata": {
        "id": "F2JUyGq4JDa8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch\n",
        "import trimesh\n",
        "import numpy as np\n",
        "from munch import *\n",
        "from options import BaseOptions\n",
        "from model import Generator\n",
        "from generate_shapes_and_images import generate\n",
        "from render_video import render_video\n",
        "\n",
        "\n",
        "torch.random.manual_seed(321)\n",
        "\n",
        "\n",
        "device = \"cuda\"\n",
        "opt = BaseOptions().parse()\n",
        "opt.camera.uniform = True\n",
        "opt.model.is_test = True\n",
        "opt.model.freeze_renderer = False\n",
        "opt.rendering.offset_sampling = True\n",
        "opt.rendering.static_viewdirs = True\n",
        "opt.rendering.force_background = True\n",
        "opt.rendering.perturb = 0\n",
        "opt.inference.renderer_output_size = opt.model.renderer_spatial_output_dim\n",
        "opt.inference.style_dim = opt.model.style_dim\n",
        "opt.inference.project_noise = opt.model.project_noise"
      ],
      "metadata": {
        "id": "Qfamt8J0JGn5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TLOiDzgKtYyi"
      },
      "source": [
        "Don't worry about this message above, \n",
        "```\n",
        "usage: ipykernel_launcher.py [-h] [--dataset_path DATASET_PATH]\n",
        "                             [--config CONFIG] [--expname EXPNAME]\n",
        "                             [--ckpt CKPT] [--continue_training]\n",
        "                             ...\n",
        "                             ...\n",
        "ipykernel_launcher.py: error: unrecognized arguments: -f /root/.local/share/jupyter/runtime/kernel-c9d47a98-bdba-4a5f-9f0a-e1437c7228b6.json\n",
        "```\n",
        "everything is perfectly fine..."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we define our model.\n",
        "\n",
        "Set the options below according to your choosing:\n",
        "1. If you plan to try the method for the AFHQ dataset (animal faces), change `model_type` to 'afhq'. Default: `ffhq` (human faces).\n",
        "2. If you wish to turn off depth rendering and marching cubes extraction and generate only RGB images, set `opt.inference.no_surface_renderings = True`. Default: `False`.\n",
        "3. If you wish to generate the image from a specific set of viewpoints, set `opt.inference.fixed_camera_angles = True`. Default: `False`.\n",
        "4. Set the number of identities you wish to create in `opt.inference.identities`. Default: `4`.\n",
        "5. Select the number of views per identity in `opt.inference.num_views_per_id`,<br>\n",
        "   (Only applicable when `opt.inference.fixed_camera_angles` is false). Default: `1`. "
      ],
      "metadata": {
        "id": "40dk1eEJo9MF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# User options\n",
        "model_type = 'ffhq' # Whether to load the FFHQ or AFHQ model\n",
        "opt.inference.no_surface_renderings = False # When true, only RGB images will be created\n",
        "opt.inference.fixed_camera_angles = False # When true, each identity will be rendered from a specific set of 13 viewpoints. Otherwise, random views are generated\n",
        "opt.inference.identities = 4 # Number of identities to generate\n",
        "opt.inference.num_views_per_id = 1 # Number of viewpoints generated per identity. This option is ignored if opt.inference.fixed_camera_angles is true.\n",
        "\n",
        "# Load saved model\n",
        "if model_type == 'ffhq':\n",
        "    model_path = 'ffhq1024x1024.pt'\n",
        "    opt.model.size = 1024\n",
        "    opt.experiment.expname = 'ffhq1024x1024'\n",
        "else:\n",
        "    opt.inference.camera.azim = 0.15\n",
        "    model_path = 'afhq512x512.pt'\n",
        "    opt.model.size = 512\n",
        "    opt.experiment.expname = 'afhq512x512'\n",
        "\n",
        "# Create results directory\n",
        "result_model_dir = 'final_model'\n",
        "results_dir_basename = os.path.join(opt.inference.results_dir, opt.experiment.expname)\n",
        "opt.inference.results_dst_dir = os.path.join(results_dir_basename, result_model_dir)\n",
        "if opt.inference.fixed_camera_angles:\n",
        "    opt.inference.results_dst_dir = os.path.join(opt.inference.results_dst_dir, 'fixed_angles')\n",
        "else:\n",
        "    opt.inference.results_dst_dir = os.path.join(opt.inference.results_dst_dir, 'random_angles')\n",
        "\n",
        "os.makedirs(opt.inference.results_dst_dir, exist_ok=True)\n",
        "os.makedirs(os.path.join(opt.inference.results_dst_dir, 'images'), exist_ok=True)\n",
        "if not opt.inference.no_surface_renderings:\n",
        "    os.makedirs(os.path.join(opt.inference.results_dst_dir, 'depth_map_meshes'), exist_ok=True)\n",
        "    os.makedirs(os.path.join(opt.inference.results_dst_dir, 'marching_cubes_meshes'), exist_ok=True)\n",
        "\n",
        "opt.inference.camera = opt.camera\n",
        "opt.inference.size = opt.model.size\n",
        "checkpoint_path = os.path.join('full_models', model_path)\n",
        "checkpoint = torch.load(checkpoint_path)\n",
        "\n",
        "# Load image generation model\n",
        "g_ema = Generator(opt.model, opt.rendering).to(device)\n",
        "pretrained_weights_dict = checkpoint[\"g_ema\"]\n",
        "model_dict = g_ema.state_dict()\n",
        "for k, v in pretrained_weights_dict.items():\n",
        "    if v.size() == model_dict[k].size():\n",
        "        model_dict[k] = v\n",
        "\n",
        "g_ema.load_state_dict(model_dict)\n",
        "\n",
        "# Load a second volume renderer that extracts surfaces at 128x128x128 (or higher) for better surface resolution\n",
        "if not opt.inference.no_surface_renderings:\n",
        "    opt['surf_extraction'] = Munch()\n",
        "    opt.surf_extraction.rendering = opt.rendering\n",
        "    opt.surf_extraction.model = opt.model.copy()\n",
        "    opt.surf_extraction.model.renderer_spatial_output_dim = 128\n",
        "    opt.surf_extraction.rendering.N_samples = opt.surf_extraction.model.renderer_spatial_output_dim\n",
        "    opt.surf_extraction.rendering.return_xyz = True\n",
        "    opt.surf_extraction.rendering.return_sdf = True\n",
        "    surface_g_ema = Generator(opt.surf_extraction.model, opt.surf_extraction.rendering, full_pipeline=False).to(device)\n",
        "\n",
        "\n",
        "    # Load weights to surface extractor\n",
        "    surface_extractor_dict = surface_g_ema.state_dict()\n",
        "    for k, v in pretrained_weights_dict.items():\n",
        "        if k in surface_extractor_dict.keys() and v.size() == surface_extractor_dict[k].size():\n",
        "            surface_extractor_dict[k] = v\n",
        "\n",
        "    surface_g_ema.load_state_dict(surface_extractor_dict)\n",
        "else:\n",
        "    surface_g_ema = None\n",
        "\n",
        "# Get the mean latent vector for g_ema\n",
        "if opt.inference.truncation_ratio < 1:\n",
        "    with torch.no_grad():\n",
        "        mean_latent = g_ema.mean_latent(opt.inference.truncation_mean, device)\n",
        "else:\n",
        "    surface_mean_latent = None\n",
        "\n",
        "# Get the mean latent vector for surface_g_ema\n",
        "if not opt.inference.no_surface_renderings:\n",
        "    surface_mean_latent = mean_latent[0]\n",
        "else:\n",
        "    surface_mean_latent = None"
      ],
      "metadata": {
        "id": "CUcWipIlpINT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generating images and meshes\n",
        "\n",
        "Finally, we run the network. The results will be saved to `evaluations/[model_name]/final_model/[fixed/random]_angles`, according to the selected setup."
      ],
      "metadata": {
        "id": "N9pqjPDYCwIJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generate(opt.inference, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent)"
      ],
      "metadata": {
        "id": "UG4hZgigDfG7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's examine the results\n",
        "\n",
        "Tip: for better mesh visualization, we recommend dowwnloading the result meshes and view them with Meshlab.\n",
        "\n",
        "Meshes loaction is: `evaluations/[model_name]/final_model/[fixed/random]_angles/[depth_map/marching_cubes]_meshes`."
      ],
      "metadata": {
        "id": "obfjXuDxNuZl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from PIL import Image\n",
        "from trimesh.viewer.notebook import scene_to_html as mesh2html\n",
        "from IPython.display import HTML as viewer_html\n",
        "\n",
        "# First let's look at the images\n",
        "img_dir = os.path.join(opt.inference.results_dst_dir,'images')\n",
        "im_list = sorted([entry for entry in os.listdir(img_dir) if 'thumb' not in entry])\n",
        "img = Image.new('RGB', (256 * len(im_list), 256))\n",
        "for i, im_file in enumerate(im_list):\n",
        "    im_path = os.path.join(img_dir, im_file)\n",
        "    curr_img = Image.open(im_path).resize((256,256)) # the displayed image is scaled to fit to the screen\n",
        "    img.paste(curr_img, (256 * i, 0))\n",
        "\n",
        "display(img)\n",
        "\n",
        "# And now, we'll move on to display the marching cubes and depth map meshes\n",
        "\n",
        "marching_cubes_meshes_dir = os.path.join(opt.inference.results_dst_dir,'marching_cubes_meshes')\n",
        "marching_cubes_meshes_list = sorted([os.path.join(marching_cubes_meshes_dir, entry) for entry in os.listdir(marching_cubes_meshes_dir) if 'obj' in entry])\n",
        "depth_map_meshes_dir = os.path.join(opt.inference.results_dst_dir,'depth_map_meshes')\n",
        "depth_map_meshes_list = sorted([os.path.join(depth_map_meshes_dir, entry) for entry in os.listdir(depth_map_meshes_dir) if 'obj' in entry])\n",
        "for i, mesh_files in enumerate(zip(marching_cubes_meshes_list, depth_map_meshes_list)):\n",
        "    mc_mesh_file, dm_mesh_file = mesh_files[0], mesh_files[1]\n",
        "    marching_cubes_mesh = trimesh.Scene(trimesh.load_mesh(mc_mesh_file, 'obj'))  \n",
        "    curr_mc_html = mesh2html(marching_cubes_mesh).replace('\"', '&quot;')\n",
        "    display(viewer_html(' '.join(['<iframe srcdoc=\"{srcdoc}\"',\n",
        "                            'width=\"{width}px\" height=\"{height}px\"',\n",
        "                            'style=\"border:none;\"></iframe>']).format(\n",
        "                            srcdoc=curr_mc_html, height=256, width=256)))\n",
        "    depth_map_mesh = trimesh.Scene(trimesh.load_mesh(dm_mesh_file, 'obj'))  \n",
        "    curr_dm_html = mesh2html(depth_map_mesh).replace('\"', '&quot;')\n",
        "    display(viewer_html(' '.join(['<iframe srcdoc=\"{srcdoc}\"',\n",
        "                            'width=\"{width}px\" height=\"{height}px\"',\n",
        "                            'style=\"border:none;\"></iframe>']).format(\n",
        "                            srcdoc=curr_dm_html, height=256, width=256)))"
      ],
      "metadata": {
        "id": "k3jXCVT7N2YZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generating videos\n",
        "\n",
        "Additionally, we can also render videos. The results will be saved to `evaluations/[model_name]/final_model/videos`.\n",
        "\n",
        "Set the options below according to your choosing:\n",
        "1. If you wish to generate only RGB videos, set `opt.inference.no_surface_renderings = True`. Default: `False`.\n",
        "2. Set the camera trajectory. To travel along the azimuth direction set `opt.inference.azim_video = True`, to travel in an ellipsoid trajectory set `opt.inference.azim_video = False`. Default: `False`.\n",
        "\n",
        "###Important Note: \n",
        " - Processing time for videos when `opt.inference.no_surface_renderings = False` is very lengthy (~ 15-20 minutes per video). Rendering each depth frame for the depth videos is very slow.<br>\n",
        " - Processing time for videos when `opt.inference.no_surface_renderings = True` is much faster (~ 1-2 minutes per video)"
      ],
      "metadata": {
        "id": "Cxq3c6anN4D0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Options\n",
        "opt.inference.no_surface_renderings = True # When true, only RGB videos will be created\n",
        "opt.inference.azim_video = True # When true, the camera trajectory will travel along the azimuth direction. Otherwise, the camera will travel along an ellipsoid trajectory.\n",
        "\n",
        "opt.inference.results_dst_dir = os.path.join(os.path.split(opt.inference.results_dst_dir)[0], 'videos')\n",
        "os.makedirs(opt.inference.results_dst_dir, exist_ok=True)\n",
        "render_video(opt.inference, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent)"
      ],
      "metadata": {
        "id": "nblhnZgcOST8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's watch the result videos.\n",
        "\n",
        "The output video files are relatively large, so it might take a while (about 1-2 minutes) for all of them to be loaded. "
      ],
      "metadata": {
        "id": "jkBYvlaeTGu1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%script bash --bg\n",
        "python3 -m https.server 8000"
      ],
      "metadata": {
        "id": "s-A9oIFXnYdM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# change ffhq1024x1024 to afhq512x512 if you are working on the AFHQ model\n",
        "%%html\n",
        "<div>\n",
        "  <video width=256 controls><source src=\"https://localhost:8000/evaluations/ffhq1024x1024/final_model/videos/sample_video_0_azim.mp4\" type=\"video/mp4\"></video>\n",
        "  <video width=256 controls><source src=\"https://localhost:8000/evaluations/ffhq1024x1024/final_model/videos/sample_video_1_azim.mp4\" type=\"video/mp4\"></video>\n",
        "  <video width=256 controls><source src=\"https://localhost:8000/evaluations/ffhq1024x1024/final_model/videos/sample_video_2_azim.mp4\" type=\"video/mp4\"></video>\n",
        "  <video width=256 controls><source src=\"https://localhost:8000/evaluations/ffhq1024x1024/final_model/videos/sample_video_3_azim.mp4\" type=\"video/mp4\"></video>\n",
        "</div>"
      ],
      "metadata": {
        "id": "Rz8tq-yYnZMD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "An alternative way to view the videos with python code. \n",
        "It loads the videos faster, but very often it crashes the notebook since the video file are too large.\n",
        "\n",
        "**It is not recommended to view the files this way**.\n",
        "\n",
        "If the notebook does crash, you can also refresh the webpage and manually download the videos.<br>\n",
        "The videos are located in `evaluations/<model_name>/final_model/videos`"
      ],
      "metadata": {
        "id": "KP_Nrl8yqL65"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from base64 import b64encode\n",
        "\n",
        "# videos_dir = opt.inference.results_dst_dir\n",
        "# videos_list = sorted([os.path.join(videos_dir, entry) for entry in os.listdir(videos_dir) if 'mp4' in entry])\n",
        "# for i, video_file in enumerate(videos_list):\n",
        "#     if i != 1:\n",
        "#         continue\n",
        "#     mp4 = open(video_file,'rb').read()\n",
        "#     data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "#     display(viewer_html(\"\"\"<video width={0} controls>\n",
        "#                                 <source src=\"{1}\" type=\"{2}\">\n",
        "#                           </video>\"\"\".format(256, data_url, \"video/mp4\")))"
      ],
      "metadata": {
        "id": "LzsQTq1hTNSH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}