differentiable_catmull_clark.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "okLalbR_g7NS"
      },
      "source": [
        "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "musUWTglgxSB"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "import torch\n",
        "need_pytorch3d=False\n",
        "try:\n",
        "    import pytorch3d\n",
        "except ModuleNotFoundError:\n",
        "    need_pytorch3d=True\n",
        "if need_pytorch3d:\n",
        "    if torch.__version__.startswith(\"1.13.\") and sys.platform.startswith(\"linux\"):\n",
        "        # We try to install PyTorch3D via a released wheel.\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 fvcore iopath\n",
        "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
        "    else:\n",
        "        # We try to install PyTorch3D from source.\n",
        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
        "        !tar xzf 1.10.0.tar.gz\n",
        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nX99zdoffBLg"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Util function for loading meshes\n",
        "from pytorch3d.io import load_objs_as_meshes, load_obj\n",
        "\n",
        "# Data structures and functions for rendering\n",
        "from pytorch3d.structures import Meshes\n",
        "from pytorch3d.vis.plotly_vis import AxisArgs, plot_batch_individually, plot_scene\n",
        "from pytorch3d.vis.texture_vis import texturesuv_image_matplotlib\n",
        "from pytorch3d.renderer import (\n",
        "    look_at_view_transform,\n",
        "    FoVPerspectiveCameras,\n",
        "    PointLights,\n",
        "    DirectionalLights,\n",
        "    Materials,\n",
        "    RasterizationSettings,\n",
        "    MeshRenderer,\n",
        "    MeshRasterizer,\n",
        "    SoftPhongShader,\n",
        "    TexturesUV,\n",
        "    TexturesVertex,\n",
        "    HardFlatShader\n",
        ")\n",
        "\n",
        "from pytorch3d.renderer import TexturesVertex\n",
        "from pytorch3d.io import load_obj, save_obj\n",
        "from pytorch3d.ops import sample_points_from_meshes\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "# add path for demo utils functions\n",
        "import sys\n",
        "import os\n",
        "sys.path.append(os.path.abspath(''))\n",
        "from pytorch3d.loss import (\n",
        "    chamfer_distance,\n",
        "    mesh_edge_loss,\n",
        "    mesh_laplacian_smoothing,\n",
        "    mesh_normal_consistency,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch3d.loss import (\n",
        "    chamfer_distance,\n",
        "    mesh_edge_loss,\n",
        "    mesh_laplacian_smoothing,\n",
        "    mesh_normal_consistency,\n",
        ")"
      ],
      "metadata": {
        "id": "rg3FxfVe8XQO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lxmehq6Zhrzv"
      },
      "source": [
        "If using **Google Colab**, fetch the utils file for plotting image grids:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HZozr3Pmho-5"
      },
      "outputs": [],
      "source": [
        "!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/docs/tutorials/utils/plot_image_grid.py\n",
        "from plot_image_grid import image_grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g4B62MzYiJUM"
      },
      "source": [
        "OR if running **locally** uncomment and run the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "paJ4Im8ahl7O"
      },
      "outputs": [],
      "source": [
        "# from utils import image_grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5jGq772XfBLk"
      },
      "source": [
        "### 1. Load a mesh and texture file\n",
        "\n",
        "Load an `.obj` file and its associated `.mtl` file and create a **Textures** and **Meshes** object.\n",
        "\n",
        "**Meshes** is a unique datastructure provided in PyTorch3D for working with batches of meshes of different sizes.\n",
        "\n",
        "**TexturesUV** is an auxiliary datastructure for storing vertex uv and texture maps for meshes.\n",
        "\n",
        "**Meshes** has several class methods which are used throughout the rendering pipeline."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a8eU4zo5jd_H"
      },
      "source": [
        "If running this notebook using **Google Colab**, run the following cell to fetch the mesh obj and texture files and save it at the path `data/cow_mesh`:\n",
        "If running locally, the data is already available at the correct path."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tTm0cVuOjb1W"
      },
      "outputs": [],
      "source": [
        "!mkdir -p data/cow_mesh\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_cube():\n",
        "  verts = np.array([\n",
        "      [-1, -1, -1],\n",
        "      [-1, -1,  1],\n",
        "      [-1,  1, -1],\n",
        "      [-1,  1,  1],\n",
        "      [ 1, -1, -1],\n",
        "      [ 1, -1,  1],\n",
        "      [ 1,  1, -1],\n",
        "      [ 1,  1,  1],\n",
        "  ], dtype=np.float32)\n",
        "\n",
        "  faces = np.array([\n",
        "      [0, 1, 3, 2],\n",
        "      [0, 2, 6, 4],\n",
        "      [0, 4, 5, 1],\n",
        "      [1, 5, 7, 3],\n",
        "      [2, 3, 7, 6],\n",
        "      [4, 6, 7, 5],\n",
        "  ])\n",
        "\n",
        "  return verts, faces"
      ],
      "metadata": {
        "id": "4HvwL3icoHQ_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def catmull_clark_subdivision(verts, faces):\n",
        "  num_verts, num_faces = verts.shape[0], faces.shape[0]\n",
        "\n",
        "  # face centroids\n",
        "  face_verts = verts[faces]\n",
        "  face_points = face_verts.mean(axis=1)\n",
        "\n",
        "  # edge data\n",
        "  # vertexIndices, adjacentFaceIndices\n",
        "  edges = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    for first, second in zip(face, face[1:]):\n",
        "\n",
        "      prev_face_idx = None\n",
        "      for edge_idx in range(len(edges)):\n",
        "        if first in edges[edge_idx][:2] and second in edges[edge_idx][:2]:\n",
        "          prev_face_idx = edges[edge_idx][-1]\n",
        "          edges[edge_idx].append(face_idx)\n",
        "          break\n",
        "\n",
        "      if prev_face_idx is None:\n",
        "        edges.append([first, second, face_idx])\n",
        "\n",
        "    prev_face_idx = None\n",
        "    for edge_idx in range(len(edges)):\n",
        "      if face[-1] in edges[edge_idx][:2] and face[0] in edges[edge_idx][:2]:\n",
        "        prev_face_idx = edges[edge_idx][-1]\n",
        "        edges[edge_idx].append(face_idx)\n",
        "        break\n",
        "\n",
        "    if prev_face_idx is None:\n",
        "      edges.append([face[-1], face[0], face_idx])\n",
        "\n",
        "  edges = np.array(edges)\n",
        "\n",
        "  edge_points = []\n",
        "  for edge in edges:\n",
        "    face_verts_mean = np.mean(face_points[edge[2:]], axis=0)\n",
        "    edge_midpoints = np.mean(verts[edge[:2]], axis=0)\n",
        "    edge_point = (edge_midpoints + face_verts_mean)/2\n",
        "    edge_points.append(edge_point)\n",
        "  edge_points = np.array(edge_points)\n",
        "\n",
        "  vert_points = []\n",
        "  for v_idx, vert in enumerate(verts):\n",
        "    vert_faces = []\n",
        "    vert_edges = []\n",
        "    vert_edge_midpoints = []\n",
        "    for face_idx, face in enumerate(faces):\n",
        "      if v_idx in face:\n",
        "        vert_faces.append(face_idx)\n",
        "    for edge_idx, edge in enumerate(edges):\n",
        "      if v_idx in edge[:2]:\n",
        "        edge_midpoint = np.mean(verts[edge[:2]], axis=0)\n",
        "        vert_edge_midpoints.append(edge_midpoint)\n",
        "        vert_edges.append(edge_idx)\n",
        "\n",
        "    F = np.mean(face_points[vert_faces], axis=0)\n",
        "    R = np.mean(vert_edge_midpoints, axis=0)\n",
        "\n",
        "    n = len(vert_faces)\n",
        "    vert_point = (F + 2*R + (n - 3)*vert) / n\n",
        "    vert_points.append(vert_point)\n",
        "\n",
        "\n",
        "  # e_idx plus offset is edge point id\n",
        "  # face_idx plus offset is face point id\n",
        "  # connections\n",
        "  res_faces = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    fp_idx = face_idx + num_verts\n",
        "\n",
        "    for e_idx, edge in enumerate(edges):\n",
        "      if face[0] in edge[:2] and face[1] in edge[:2]:\n",
        "        ep1 = e_idx + num_verts + num_faces\n",
        "      elif face[1] in edge[:2] and face[2] in edge[:2]:\n",
        "        ep2 = e_idx + num_verts + num_faces\n",
        "      elif face[2] in edge[:2] and face[3] in edge[:2]:\n",
        "        ep3 = e_idx + num_verts + num_faces\n",
        "      elif face[3] in edge[:2] and face[0] in edge[:2]:\n",
        "        ep4 = e_idx + num_verts + num_faces\n",
        "    new_face1 = [face[0], ep1, fp_idx, ep4]\n",
        "    new_face2 = [ep1, face[1], ep2, fp_idx]\n",
        "    new_face3 = [fp_idx, ep2, face[2], ep3]\n",
        "    new_face4 = [ep4, fp_idx, ep3, face[3]]\n",
        "    res_faces.extend([new_face1, new_face2, new_face3, new_face4])\n",
        "\n",
        "\n",
        "  res_verts = np.concatenate((vert_points, face_points, edge_points), axis=0)\n",
        "  res_faces = np.array(res_faces)\n",
        "\n",
        "  return res_verts, res_faces"
      ],
      "metadata": {
        "id": "3dunl2VDmfnN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def quads_to_tris(faces):\n",
        "  new_faces = []\n",
        "  for face in faces:\n",
        "    f1 = face[:-1]\n",
        "    f2 = [face[-2], face[-1], face[0]]\n",
        "    new_faces.append(f1)\n",
        "    new_faces.append(f2)\n",
        "\n",
        "  return np.array(new_faces)"
      ],
      "metadata": {
        "id": "GPvkVFTepB3M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def catmull_clark_subdivision_pyt(verts, faces):\n",
        "  num_verts, num_faces = verts.shape[0], faces.shape[0]\n",
        "\n",
        "  # face centroids\n",
        "  face_verts = verts[faces]\n",
        "  face_points = face_verts.mean(axis=1)\n",
        "\n",
        "  # edge data\n",
        "  # vertexIndices, adjacentFaceIndices\n",
        "  edges = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    for first, second in zip(face, face[1:]):\n",
        "\n",
        "      prev_face_idx = None\n",
        "      for edge_idx in range(len(edges)):\n",
        "        if first in edges[edge_idx][:2] and second in edges[edge_idx][:2]:\n",
        "          prev_face_idx = edges[edge_idx][-1]\n",
        "          edges[edge_idx].append(face_idx)\n",
        "          break\n",
        "\n",
        "      if prev_face_idx is None:\n",
        "        edges.append([first, second, face_idx])\n",
        "\n",
        "    prev_face_idx = None\n",
        "    for edge_idx in range(len(edges)):\n",
        "      if face[-1] in edges[edge_idx][:2] and face[0] in edges[edge_idx][:2]:\n",
        "        prev_face_idx = edges[edge_idx][-1]\n",
        "        edges[edge_idx].append(face_idx)\n",
        "        break\n",
        "\n",
        "    if prev_face_idx is None:\n",
        "      edges.append([face[-1], face[0], face_idx])\n",
        "\n",
        "  edges = torch.tensor(edges)\n",
        "\n",
        "  edge_points = []\n",
        "  for edge in edges:\n",
        "    face_verts_mean = torch.mean(face_points[edge[2:]], axis=0)\n",
        "    edge_midpoints = torch.mean(verts[edge[:2]], axis=0)\n",
        "    edge_point = (edge_midpoints + face_verts_mean)/2\n",
        "    edge_points.append(edge_point)\n",
        "  edge_points = torch.stack(edge_points)\n",
        "\n",
        "  vert_points = []\n",
        "  for v_idx, vert in enumerate(verts):\n",
        "    vert_faces = []\n",
        "    vert_edges = []\n",
        "    vert_edge_midpoints = []\n",
        "    for face_idx, face in enumerate(faces):\n",
        "      if v_idx in face:\n",
        "        vert_faces.append(face_idx)\n",
        "    for edge_idx, edge in enumerate(edges):\n",
        "      if v_idx in edge[:2]:\n",
        "        edge_midpoint = torch.mean(verts[edge[:2]], axis=0)\n",
        "        vert_edge_midpoints.append(edge_midpoint)\n",
        "        vert_edges.append(edge_idx)\n",
        "\n",
        "    F = torch.mean(face_points[vert_faces], axis=0)\n",
        "    R = torch.mean(torch.stack(vert_edge_midpoints), axis=0)\n",
        "\n",
        "    n = len(vert_faces)\n",
        "    vert_point = (F + 2*R + (n - 3)*vert) / n\n",
        "    vert_points.append(vert_point)\n",
        "\n",
        "\n",
        "  # e_idx plus offset is edge point id\n",
        "  # face_idx plus offset is face point id\n",
        "  # connections\n",
        "  res_faces = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    fp_idx = face_idx + num_verts\n",
        "\n",
        "    for e_idx, edge in enumerate(edges):\n",
        "      if face[0] in edge[:2] and face[1] in edge[:2]:\n",
        "        ep1 = e_idx + num_verts + num_faces\n",
        "      elif face[1] in edge[:2] and face[2] in edge[:2]:\n",
        "        ep2 = e_idx + num_verts + num_faces\n",
        "      elif face[2] in edge[:2] and face[3] in edge[:2]:\n",
        "        ep3 = e_idx + num_verts + num_faces\n",
        "      elif face[3] in edge[:2] and face[0] in edge[:2]:\n",
        "        ep4 = e_idx + num_verts + num_faces\n",
        "    new_face1 = [face[0], ep1, fp_idx, ep4]\n",
        "    new_face2 = [ep1, face[1], ep2, fp_idx]\n",
        "    new_face3 = [fp_idx, ep2, face[2], ep3]\n",
        "    new_face4 = [ep4, fp_idx, ep3, face[3]]\n",
        "    res_faces.extend([new_face1, new_face2, new_face3, new_face4])\n",
        "\n",
        "\n",
        "  res_verts = torch.concatenate((torch.stack(vert_points), face_points, edge_points), axis=0)\n",
        "  res_faces = torch.tensor(res_faces)\n",
        "\n",
        "  return res_verts, res_faces"
      ],
      "metadata": {
        "id": "bmrnVwWNYplj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def quads_to_tris_pyt(faces):\n",
        "  new_faces = []\n",
        "  for face in faces:\n",
        "    f1 = face[:-1]\n",
        "    f2 = torch.tensor([face[-2], face[-1], face[0]])\n",
        "    new_faces.append(f1)\n",
        "    new_faces.append(f2)\n",
        "\n",
        "  return torch.stack(new_faces)"
      ],
      "metadata": {
        "id": "NQqTQ5sbe_is"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "verts, faces = get_cube()\n",
        "verts = torch.tensor(verts)\n",
        "faces = torch.tensor(faces)\n",
        "verts, faces = catmull_clark_subdivision_pyt(verts, faces)\n",
        "faces = quads_to_tris_pyt(faces)"
      ],
      "metadata": {
        "id": "Tfs-n9APosV9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "device = torch.device('cuda')\n",
        "\n",
        "verts = torch.tensor(verts).to(device)\n",
        "faces = torch.tensor(faces).to(device)\n",
        "\n",
        "#final_obj = os.path.join('./', 'cube1.obj')\n",
        "#save_obj(final_obj, verts, faces)"
      ],
      "metadata": {
        "id": "YXSGgUQi1YEY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "verts_rgb = torch.ones_like(verts)[None] # (1, V, 3)\n",
        "textures = TexturesVertex(verts_features=verts_rgb.to(device))\n",
        "mesh = Meshes(verts=[verts], faces=[faces], textures=textures)"
      ],
      "metadata": {
        "id": "2WeE85v3k_y9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gt_verts, gt_faces, gt_aux = load_obj('gt_mesh.obj')\n",
        "\n",
        "gt_verts = gt_verts.to(device)\n",
        "gt_faces = gt_faces.verts_idx.to(device)\n",
        "\n",
        "gt_mesh = Meshes(verts=[gt_verts], faces=[gt_faces])\n"
      ],
      "metadata": {
        "id": "4EYigdxg5_t_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig2 = plot_scene({\n",
        "    \"plot1\": {\n",
        "        \"gt_mesh\": gt_mesh\n",
        "    }\n",
        "},\n",
        "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
        "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
        "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"},\n",
        "    axis_args=AxisArgs(showgrid=True))\n",
        "fig2.show()"
      ],
      "metadata": {
        "id": "vBQexhTR9OFJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "verts_cage, faces_cage = get_cube()\n",
        "verts_subd, faces_subd = catmull_clark_subdivision(verts_cage, faces_cage)\n",
        "\n",
        "verts = verts_subd\n",
        "faces = faces_subd\n",
        "\n",
        "verts = torch.tensor(verts).to(device)\n",
        "faces = torch.tensor(quads_to_tris(faces)).to(device)\n",
        "\n",
        "\n",
        "mesh = Meshes(verts=[verts], faces=[faces])\n",
        "\n",
        "\n",
        "\n",
        "src_pts = sample_points_from_meshes(mesh, 1500)\n",
        "gt_pts = sample_points_from_meshes(gt_mesh, 1500)\n",
        "loss_chamfer, _ = chamfer_distance(src_pts, gt_pts)"
      ],
      "metadata": {
        "id": "F7feej0W4Vke"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_cube_pyt():\n",
        "  verts = torch.tensor([\n",
        "      [-1, -1, -1],\n",
        "      [-1, -1,  1],\n",
        "      [-1,  1, -1],\n",
        "      [-1,  1,  1],\n",
        "      [ 1, -1, -1],\n",
        "      [ 1, -1,  1],\n",
        "      [ 1,  1, -1],\n",
        "      [ 1,  1,  1],\n",
        "  ], dtype=torch.float)\n",
        "\n",
        "  faces = torch.tensor([\n",
        "      [0, 1, 3, 2],\n",
        "      [0, 2, 6, 4],\n",
        "      [0, 4, 5, 1],\n",
        "      [1, 5, 7, 3],\n",
        "      [2, 3, 7, 6],\n",
        "      [4, 6, 7, 5],\n",
        "  ])\n",
        "\n",
        "  return verts, faces"
      ],
      "metadata": {
        "id": "JLjHAIbcf1Pc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v, f = get_cube_pyt()"
      ],
      "metadata": {
        "id": "mK6Xs8kYFcs9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# run multiple times to restore multiple levels\n",
        "verts_cage, faces_cage = v, f\n",
        "verts_cage = torch.autograd.Variable(verts_cage.data, requires_grad=True)\n",
        "optimizer = torch.optim.SGD([verts_cage], lr=0.15, momentum=0.9)\n",
        "\n",
        "iter = range(100)\n",
        "for i in iter:\n",
        "  optimizer.zero_grad()\n",
        "\n",
        "  v,f = verts_cage, faces_cage\n",
        "  subd_count = 1\n",
        "  for j in range(subd_count):\n",
        "    v, f = catmull_clark_subdivision_pyt(v, f)\n",
        "\n",
        "\n",
        "  mesh = Meshes(verts=[v], faces=[quads_to_tris_pyt(f)])\n",
        "\n",
        "\n",
        "  pt_num = 15000\n",
        "  src_pts = sample_points_from_meshes(mesh, pt_num).to(device)\n",
        "  gt_pts = sample_points_from_meshes(gt_mesh, pt_num).to(device)\n",
        "\n",
        "\n",
        "  loss_chamfer, _ = chamfer_distance(src_pts, gt_pts)\n",
        "  loss_edge = mesh_edge_loss(mesh)\n",
        "  loss_normal = mesh_normal_consistency(mesh)\n",
        "  loss_laplacian = mesh_laplacian_smoothing(mesh, method=\"uniform\")\n",
        "\n",
        "  loss = loss_chamfer + loss_edge + loss_normal + loss_laplacian\n",
        "\n",
        "  if i % 5 == 0:\n",
        "    print(f'i {i} - {loss}')\n",
        "\n",
        "  loss.backward(retain_graph=True)\n",
        "  optimizer.step()"
      ],
      "metadata": {
        "id": "5SmS285uFqvs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(verts_cage)"
      ],
      "metadata": {
        "id": "x_mKfsT1DIvH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tYs7jCeRwHkh"
      },
      "outputs": [],
      "source": [
        "v,f = verts_cage, faces_cage\n",
        "subd_level = 1\n",
        "for j in range(subd_level):\n",
        "  v, f = catmull_clark_subdivision_pyt(v, f)\n",
        "\n",
        "mesh = Meshes(verts=[v], faces=[quads_to_tris_pyt(f)])\n",
        "\n",
        "fig2 = plot_scene({\n",
        "    \"cow_plot1\": {\n",
        "        \"cows\": mesh\n",
        "    }\n",
        "},\n",
        "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
        "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
        "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"},\n",
        "    axis_args=AxisArgs(showgrid=True))\n",
        "fig2.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v,f = verts_cage, faces_cage\n",
        "with open('export.obj', 'w') as file:\n",
        "    file.write(\"# OBJ file\\n\")\n",
        "    for vertex in v:\n",
        "        file.write(f\"v {vertex[0]} {vertex[1]} {vertex[2]}\\n\")\n",
        "    for face in f:\n",
        "        file.write(\"f\")\n",
        "        for i in face:\n",
        "            file.write(\" %d\" % (i + 1))\n",
        "        file.write(\"\\n\")"
      ],
      "metadata": {
        "id": "X-Z47gO2dljk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "verts_cage, faces_cage = get_cube_pyt()\n",
        "mesh = Meshes(verts=[verts_cage], faces=[quads_to_tris_pyt(faces_cage)])\n",
        "\n",
        "fig2 = plot_scene({\n",
        "    \"cow_plot1\": {\n",
        "        \"cows\": gt_mesh\n",
        "    }\n",
        "},\n",
        "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
        "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
        "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"},\n",
        "    axis_args=AxisArgs(showgrid=True))\n",
        "fig2.show()"
      ],
      "metadata": {
        "id": "swYkxzUMidzn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "def get_cube_pyt():\n",
        "  verts = torch.tensor([\n",
        "      [-1, -1, -1],\n",
        "      [-1, -1,  1],\n",
        "      [-1,  1, -1],\n",
        "      [-1,  1,  1],\n",
        "      [ 1, -1, -1],\n",
        "      [ 1, -1,  1],\n",
        "      [ 1,  1, -1],\n",
        "      [ 1,  1,  1],\n",
        "  ], dtype=torch.float)\n",
        "\n",
        "  faces = torch.tensor([\n",
        "      [0, 1, 3, 2],\n",
        "      [0, 2, 6, 4],\n",
        "      [0, 4, 5, 1],\n",
        "      [1, 5, 7, 3],\n",
        "      [2, 3, 7, 6],\n",
        "      [4, 6, 7, 5],\n",
        "  ])\n",
        "\n",
        "  return verts, faces\n",
        "\n",
        "def catmull_clark_subdivision(verts, faces):\n",
        "  num_verts, num_faces = verts.shape[0], faces.shape[0]\n",
        "\n",
        "  # face centroids\n",
        "  face_verts = verts[faces]\n",
        "  face_points = face_verts.mean(axis=1)\n",
        "\n",
        "  # edge data\n",
        "  # vertexIndices, adjacentFaceIndices\n",
        "  edges = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    for first, second in zip(face, face[1:]):\n",
        "\n",
        "      prev_face_idx = None\n",
        "      for edge_idx in range(len(edges)):\n",
        "        if first in edges[edge_idx][:2] and second in edges[edge_idx][:2]:\n",
        "          prev_face_idx = edges[edge_idx][-1]\n",
        "          edges[edge_idx].append(face_idx)\n",
        "          break\n",
        "\n",
        "      if prev_face_idx is None:\n",
        "        edges.append([first, second, face_idx])\n",
        "\n",
        "    prev_face_idx = None\n",
        "    for edge_idx in range(len(edges)):\n",
        "      if face[-1] in edges[edge_idx][:2] and face[0] in edges[edge_idx][:2]:\n",
        "        prev_face_idx = edges[edge_idx][-1]\n",
        "        edges[edge_idx].append(face_idx)\n",
        "        break\n",
        "\n",
        "    if prev_face_idx is None:\n",
        "      edges.append([face[-1], face[0], face_idx])\n",
        "\n",
        "  edges = np.array(edges)\n",
        "\n",
        "  edge_points = []\n",
        "  for edge in edges:\n",
        "    face_verts_mean = np.mean(face_points[edge[2:]], axis=0)\n",
        "    edge_midpoints = np.mean(verts[edge[:2]], axis=0)\n",
        "    edge_point = (edge_midpoints + face_verts_mean)/2\n",
        "    edge_points.append(edge_point)\n",
        "  edge_points = np.array(edge_points)\n",
        "\n",
        "  vert_points = []\n",
        "  for v_idx, vert in enumerate(verts):\n",
        "    vert_faces = []\n",
        "    vert_edges = []\n",
        "    vert_edge_midpoints = []\n",
        "    for face_idx, face in enumerate(faces):\n",
        "      if v_idx in face:\n",
        "        vert_faces.append(face_idx)\n",
        "    for edge_idx, edge in enumerate(edges):\n",
        "      if v_idx in edge[:2]:\n",
        "        edge_midpoint = np.mean(verts[edge[:2]], axis=0)\n",
        "        vert_edge_midpoints.append(edge_midpoint)\n",
        "        vert_edges.append(edge_idx)\n",
        "\n",
        "    F = np.mean(face_points[vert_faces], axis=0)\n",
        "    R = np.mean(vert_edge_midpoints, axis=0)\n",
        "\n",
        "    n = len(vert_faces)\n",
        "    vert_point = (F + 2*R + (n - 3)*vert) / n\n",
        "    vert_points.append(vert_point)\n",
        "\n",
        "\n",
        "  # e_idx plus offset is edge point id\n",
        "  # face_idx plus offset is face point id\n",
        "  # connections\n",
        "  res_faces = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    fp_idx = face_idx + num_verts\n",
        "\n",
        "    for e_idx, edge in enumerate(edges):\n",
        "      if face[0] in edge[:2] and face[1] in edge[:2]:\n",
        "        ep1 = e_idx + num_verts + num_faces\n",
        "      elif face[1] in edge[:2] and face[2] in edge[:2]:\n",
        "        ep2 = e_idx + num_verts + num_faces\n",
        "      elif face[2] in edge[:2] and face[3] in edge[:2]:\n",
        "        ep3 = e_idx + num_verts + num_faces\n",
        "      elif face[3] in edge[:2] and face[0] in edge[:2]:\n",
        "        ep4 = e_idx + num_verts + num_faces\n",
        "    new_face1 = [face[0], ep1, fp_idx, ep4]\n",
        "    new_face2 = [ep1, face[1], ep2, fp_idx]\n",
        "    new_face3 = [fp_idx, ep2, face[2], ep3]\n",
        "    new_face4 = [ep4, fp_idx, ep3, face[3]]\n",
        "    res_faces.extend([new_face1, new_face2, new_face3, new_face4])\n",
        "\n",
        "\n",
        "  res_verts = np.concatenate((vert_points, face_points, edge_points), axis=0)\n",
        "  res_faces = np.array(res_faces)\n",
        "\n",
        "  return res_verts, res_faces\n",
        "\n",
        "def catmull_clark_subdivision_pyt(verts, faces):\n",
        "  num_verts, num_faces = verts.shape[0], faces.shape[0]\n",
        "\n",
        "  # face centroids\n",
        "  face_verts = verts[faces]\n",
        "  face_points = face_verts.mean(dim=1)\n",
        "\n",
        "  # edge data\n",
        "  # vertexIndices, adjacentFaceIndices\n",
        "  edges = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    for first, second in zip(face, face[1:]):\n",
        "\n",
        "      prev_face_idx = None\n",
        "      for edge_idx in range(len(edges)):\n",
        "        if first in edges[edge_idx][:2] and second in edges[edge_idx][:2]:\n",
        "          prev_face_idx = edges[edge_idx][-1]\n",
        "          edges[edge_idx].append(face_idx)\n",
        "          break\n",
        "\n",
        "      if prev_face_idx is None:\n",
        "        edges.append([first, second, face_idx])\n",
        "\n",
        "    prev_face_idx = None\n",
        "    for edge_idx in range(len(edges)):\n",
        "      if face[-1] in edges[edge_idx][:2] and face[0] in edges[edge_idx][:2]:\n",
        "        prev_face_idx = edges[edge_idx][-1]\n",
        "        edges[edge_idx].append(face_idx)\n",
        "        break\n",
        "\n",
        "    if prev_face_idx is None:\n",
        "      edges.append([face[-1], face[0], face_idx])\n",
        "\n",
        "  edges = torch.tensor(edges)\n",
        "\n",
        "  edge_points = []\n",
        "  for edge in edges:\n",
        "    face_verts_mean = torch.mean(face_points[edge[2:]], axis=0)\n",
        "    edge_midpoints = torch.mean(verts[edge[:2]], axis=0)\n",
        "    edge_point = (edge_midpoints + face_verts_mean)/2\n",
        "    edge_points.append(edge_point)\n",
        "  edge_points = torch.stack(edge_points)\n",
        "\n",
        "  vert_points = []\n",
        "  for v_idx, vert in enumerate(verts):\n",
        "    vert_faces = []\n",
        "    vert_edges = []\n",
        "    vert_edge_midpoints = []\n",
        "    for face_idx, face in enumerate(faces):\n",
        "      if v_idx in face:\n",
        "        vert_faces.append(face_idx)\n",
        "    for edge_idx, edge in enumerate(edges):\n",
        "      if v_idx in edge[:2]:\n",
        "        edge_midpoint = torch.mean(verts[edge[:2]], axis=0)\n",
        "        vert_edge_midpoints.append(edge_midpoint)\n",
        "        vert_edges.append(edge_idx)\n",
        "\n",
        "    F = torch.mean(face_points[vert_faces], axis=0)\n",
        "    R = torch.mean(torch.stack(vert_edge_midpoints), axis=0)\n",
        "\n",
        "    n = len(vert_faces)\n",
        "    vert_point = (F + 2*R + (n - 3)*vert) / n\n",
        "    vert_points.append(vert_point)\n",
        "\n",
        "\n",
        "  # e_idx plus offset is edge point id\n",
        "  # face_idx plus offset is face point id\n",
        "  # connections\n",
        "  res_faces = []\n",
        "  for face_idx, face in enumerate(faces):\n",
        "    fp_idx = face_idx + num_verts\n",
        "\n",
        "    for e_idx, edge in enumerate(edges):\n",
        "      if face[0] in edge[:2] and face[1] in edge[:2]:\n",
        "        ep1 = e_idx + num_verts + num_faces\n",
        "      elif face[1] in edge[:2] and face[2] in edge[:2]:\n",
        "        ep2 = e_idx + num_verts + num_faces\n",
        "      elif face[2] in edge[:2] and face[3] in edge[:2]:\n",
        "        ep3 = e_idx + num_verts + num_faces\n",
        "      elif face[3] in edge[:2] and face[0] in edge[:2]:\n",
        "        ep4 = e_idx + num_verts + num_faces\n",
        "    new_face1 = [face[0], ep1, fp_idx, ep4]\n",
        "    new_face2 = [ep1, face[1], ep2, fp_idx]\n",
        "    new_face3 = [fp_idx, ep2, face[2], ep3]\n",
        "    new_face4 = [ep4, fp_idx, ep3, face[3]]\n",
        "    res_faces.extend([new_face1, new_face2, new_face3, new_face4])\n",
        "\n",
        "\n",
        "  res_verts = torch.concatenate((torch.stack(vert_points), face_points, edge_points), axis=0)\n",
        "  res_faces = torch.tensor(res_faces)\n",
        "\n",
        "  return res_verts, res_faces\n",
        "\n",
        "v, f = get_cube_pyt()\n",
        "v1, f1 = catmull_clark_subdivision(v.numpy(),f.numpy())\n",
        "#print(v1,f1)\n",
        "v, f = catmull_clark_subdivision_pyt(v,f)\n",
        "test1 = (v.numpy()==v1).all()\n",
        "test2 = (f.numpy()==f1).all()\n",
        "print(test1)\n",
        "print(test2)\n"
      ],
      "metadata": {
        "id": "Ycv14zKfnT_S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "start_time = time.time()\n",
        "v, f = get_cube_pyt()\n",
        "for i in range(4):\n",
        "  v, f = catmull_clark_subdivision_pyt(v,f)\n",
        "print(\"--- %s seconds ---\" % (time.time() - start_time))"
      ],
      "metadata": {
        "id": "Xxd725m3o4EU"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "anp_metadata": {
      "path": "notebooks/render_textured_meshes.ipynb"
    },
    "bento_stylesheets": {
      "bento/extensions/flow/main.css": true,
      "bento/extensions/kernel_selector/main.css": true,
      "bento/extensions/kernel_ui/main.css": true,
      "bento/extensions/new_kernel/main.css": true,
      "bento/extensions/system_usage/main.css": true,
      "bento/extensions/theme/main.css": true
    },
    "colab": {
      "provenance": []
    },
    "disseminate_notebook_info": {
      "backup_notebook_id": "569222367081034"
    },
    "kernelspec": {
      "display_name": "pytorch3d_etc (local)",
      "language": "python",
      "name": "pytorch3d_etc_local"
    },
    "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.7.5+"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}