PanoHead_colab_basic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "A100",
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyMQnZoDVwREugk5QHqJIkuV",
      "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/hack-mans/PanoHead/blob/main/PanoHead_colab_basic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 00 - Check GPU status (requires 8GB VRAM)\n",
        "\n"
      ],
      "metadata": {
        "id": "eNxbH1ADnoo5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UoGA3xmna5P"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 01 - Setup"
      ],
      "metadata": {
        "id": "vBxQheQ4o6ea"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Clone the PanoHead repo"
      ],
      "metadata": {
        "id": "y35sdV1ZoaVG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/hack-mans/PanoHead.git\n",
        "%cd PanoHead"
      ],
      "metadata": {
        "id": "Jzs34AIooh7x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install requirements"
      ],
      "metadata": {
        "id": "epR6RelHomD4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install imgui glfw pyspng mrcfile ninja plyfile trimesh"
      ],
      "metadata": {
        "id": "FbN7YS8KpHEZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Fetch the models"
      ],
      "metadata": {
        "id": "r6LFS4YDMezZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir '/content/PanoHead/models'\n",
        "!gdown 1FqvQzICV1H4fbQaz8BiWxtiRYxJd4T8N -O '/content/PanoHead/models/easy-khair-180-gpc0.8-trans10-025000.pkl'\n",
        "!gdown 1bYg9jCM9KJB2qxhFpktUDL_UwsaaitSS -O '/content/PanoHead/models/baseline-easy-khair-025000.pkl'\n",
        "!gdown 1Qb4vD7Ylbt9BLHVqa2uljvdy0jNRB1hd -O '/content/PanoHead/models/ablation-trigridD-1-025000.pkl'"
      ],
      "metadata": {
        "id": "PwwMqLj5-vvY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 02 - Setup Datasets (unused currently - WIP)"
      ],
      "metadata": {
        "id": "nAOFAoqLQYm8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Upload your images to /content/datasets/PanoHead/datasets/custom_img"
      ],
      "metadata": {
        "id": "EReACGrlQbsR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir /content/PanoHead/datasets/custom_img\n",
        "import os\n",
        "os.chdir('content/PanoHead/datasets/custom_img')\n",
        "from google.colab import files\n",
        "uploaded - files.upload()"
      ],
      "metadata": {
        "id": "37hFPOEiQpav"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Upload your segmentation masks to /content/PanoHead/datasets/custom_seg"
      ],
      "metadata": {
        "id": "ZV8M-A9zAckT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir /content/PanoHead/datasets/custom_seg\n",
        "import os\n",
        "os.chdir('content/PanoHead/datasets/custom_seg')\n",
        "from google.colab import files\n",
        "uploaded - files.upload()"
      ],
      "metadata": {
        "id": "X8D8Liv-AstB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 03A - Generating Results (optional)"
      ],
      "metadata": {
        "id": "co1tArQpMlMl"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate videos using pre-trained model"
      ],
      "metadata": {
        "id": "BJmaJg2MMqbP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python gen_videos.py --network models/easy-khair-180-gpc0.8-trans10-025000.pkl --seeds 0-3 --grid 2x2 --outdir=out --cfg Head --trunc 0.7"
      ],
      "metadata": {
        "id": "nLGlZu7-M0Rr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 03B - Applications"
      ],
      "metadata": {
        "id": "q90d_aeFNCX9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate full head reconstruction 360 videos and PLY 3D mesh from single RGB image\n",
        "(using easy-khair-180-gpc0.8-trans10-025000.pkl)"
      ],
      "metadata": {
        "id": "8EmtM-hDNN5T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/PanoHead\n",
        "# Perform the pti and save w\n",
        "!python projector_withseg.py --outdir=\"/content/PanoHead/pti_out\" --target_img=\"/content/PanoHead/dataset/testdata_img\" --target_seg=\"/content/PanoHead/dataset/testdata_seg\" --network \"/content/PanoHead/models/easy-khair-180-gpc0.8-trans10-025000.pkl\" --idx 0\n",
        "# Generate .mp4 before finetune\n",
        "!python gen_videos_proj_withseg.py --output=\"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/pre.mp4\" --latent=\"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/projected_w.npz\" --trunc 0.7 --network \"/content/PanoHead/models/easy-khair-180-gpc0.8-trans10-025000.pkl\" --cfg Head\n",
        "# Generate .mp4, .ply mesh and frame images after finetune\n",
        "!python gen_videos_proj_withseg.py --output=\"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post.mp4\" --latent=\"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/projected_w.npz\" --trunc 0.7 --network \"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/fintuned_generator.pkl\" --cfg Head --shapes True --frames True --level 42"
      ],
      "metadata": {
        "id": "HoCcz0ZPMJBq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 04 Display Results"
      ],
      "metadata": {
        "id": "iYIFGIQtP2MJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "import os\n",
        "\n",
        "# Input video path\n",
        "video_path = '/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post.mp4'\n",
        "\n",
        "# Compressed video path\n",
        "compressed_path = \"/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post_compressed.mp4\"\n",
        "\n",
        "os.system(f\"ffmpeg -i {video_path} -vcodec libx264 {compressed_path}\")\n",
        "\n",
        "# Show video\n",
        "mp4 = open(compressed_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "HTML(\"\"\"\n",
        "<video width=512 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ],
      "metadata": {
        "id": "Jgmz3-nGP5L7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 05 Download Results"
      ],
      "metadata": {
        "id": "hLvhH0FwMCxw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import shutil\n",
        "from google.colab import files\n",
        "\n",
        "def zip_folder(folder_path, output_path):\n",
        "  shutil.make_archive(output_path, 'zip', folder_path)\n",
        "\n",
        "# Zip the results folder\n",
        "zip_folder('/content/PanoHead/pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render', '/content/panohead_results')\n",
        "\n",
        "# Download the zip file\n",
        "files.download('/content/panohead_results.zip')"
      ],
      "metadata": {
        "id": "lNiuQpLoMGPR"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}