Automate HairStep Inference with Model Download, Setup, and Execution

notebooks/HairStep_Inference.ipynb

@@ -0,0 +1,178 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import shutil\n",
+    "import os\n",
+    "import subprocess"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Configs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAM_MODEL_URL = \"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\"\n",
+    "OCC_ORIENT_MODELS_URL = \"https://drive.google.com/uc?id=1-akuukaYYtJDta24AAqVdgUOGte4EmQf\"\n",
+    "\n",
+    "SAM_MODEL_DIR = \"./checkpoints/SAM-models/\"\n",
+    "OCC_ORIENT_MODEL_DIR = \"./checkpoints/recon3D/\"\n",
+    "\n",
+    "SAM_MODEL_PATH = os.path.join(SAM_MODEL_DIR, \"sam_vit_h_4b8939.pth\")\n",
+    "OCC_ORIENT_MODEL_PATH = os.path.join(OCC_ORIENT_MODEL_DIR, \"recon3D.zip\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Setup the Project"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Download Pretained Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "SAM model already exists\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download SAM\n",
+    "if os.path.exists(SAM_MODEL_PATH):\n",
+    "    print(\"SAM model already exists\")\n",
+    "else:\n",
+    "    os.makedirs(SAM_MODEL_DIR, exist_ok=True)\n",
+    "    subprocess.run([\"curl\", \"-L\", SAM_MODEL_URL, \"-o\", SAM_MODEL_PATH])\n",
+    "    print(\"SAM model downloaded successfully\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "OCC and ORIENT models already exists\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download OCC and ORIENT models\n",
+    "if os.path.exists(OCC_ORIENT_MODEL_PATH):\n",
+    "    print(\"OCC and ORIENT models already exists\")\n",
+    "else:\n",
+    "    import gdown\n",
+    "    os.makedirs(OCC_ORIENT_MODEL_DIR, exist_ok=True)\n",
+    "    gdown.download(OCC_ORIENT_MODELS_URL, OCC_ORIENT_MODEL_PATH, quiet=False)\n",
+    "    shutil.unpack_archive(OCC_ORIENT_MODEL_PATH, OCC_ORIENT_MODEL_DIR)\n",
+    "    print(\"OCC and ORIENT models downloaded and extracted successfully\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Build External Modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Change directory to external/3DDFA_V2\n",
+    "os.chdir(\"external/3DDFA_V2\")\n",
+    "\n",
+    "# Execute the build script\n",
+    "subprocess.run(\n",
+    "    [\"sh\", \"./build.sh\"],\n",
+    "    stdout=subprocess.PIPE, stderr=subprocess.STDOUT,\n",
+    ")\n",
+    "\n",
+    "# Change directory back to the original location\n",
+    "os.chdir(\"../../\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Execute Scripts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "commands = [\n",
+    "    \"CUDA_VISIBLE_DEVICES=0 python -m scripts.img2hairstep\",\n",
+    "    \"CUDA_VISIBLE_DEVICES=0 python scripts/get_lmk.py\",\n",
+    "    \"CUDA_VISIBLE_DEVICES=0 python -m scripts.opt_cam\",\n",
+    "    \"CUDA_VISIBLE_DEVICES=0 python -m scripts.recon3D\"\n",
+    "]\n",
+    "\n",
+    "for command in commands:\n",
+    "    subprocess.run(command, shell=True)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "clip",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}