visualization_scripts/mesh_visualizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6a29b9c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob, math, random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import matplotlib\n",
    "import trimesh\n",
    "import plyfile\n",
    "import open3d as o3d\n",
    "from open3d_utils import make_pcd, make_line_set, make_mesh\n",
    "import pyvista as pv\n",
    "from pyvista import examples\n",
    "pv.rcParams['transparent_background'] = True\n",
    "\n",
    "# export DISPLAY for remote visualization\n",
    "# os.environ['DISPLAY'] = \":77\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b5f4a2d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_plydata(mesh_ply_file):\n",
    "    plydata = plyfile.PlyData.read(mesh_ply_file)\n",
    "    mesh_faces = np.stack([list(item) for item in plydata['face']])[:,0]\n",
    "\n",
    "    plydata_vertices_data = np.asarray(list(plydata['vertex']))\n",
    "    mesh_vertices_data = []\n",
    "    for line in plydata_vertices_data:\n",
    "        mesh_vertices_data.append(np.asarray(list(line)))\n",
    "    mesh_vertices_data = np.stack(mesh_vertices_data)\n",
    "    mesh_vertices = mesh_vertices_data[:,:3]\n",
    "    mesh_colors = mesh_vertices_data[:,3:]\n",
    "    mesh_colors = (mesh_colors - mesh_colors.min(axis=0, keepdims=True)) / (mesh_colors.max(axis=0, keepdims=True) - mesh_colors.min(axis=0, keepdims=True))\n",
    "    return mesh_vertices, mesh_faces, mesh_colors\n",
    "        \n",
    "\n",
    "def visualize_mesh(mesh_ply_file, deformations_as_color=False):\n",
    "    if deformations_as_color:\n",
    "        mesh_vertices, mesh_faces, mesh_colors = read_plydata(mesh_ply_file)\n",
    "        o3d.visualization.draw_geometries([make_mesh(mesh_vertices, mesh_faces, mesh_colors)])\n",
    "    else:\n",
    "        mesh = o3d.io.read_triangle_mesh(mesh_ply_file)\n",
    "        o3d.visualization.draw_geometries([make_mesh(mesh.vertices, mesh.triangles)])\n",
    "        \n",
    "    \n",
    "def render_mesh_gif(mesh_ply_file, mesh_save_path, deformations_as_color=False):\n",
    "    if deformations_as_color:\n",
    "        mesh_vertices, mesh_faces, mesh_colors = read_plydata(mesh_ply_file)\n",
    "    else:\n",
    "        mesh = o3d.io.read_triangle_mesh(mesh_ply_file)\n",
    "        mesh_vertices, mesh_faces, mesh_colors = np.asarray(mesh.vertices), np.asarray(mesh.triangles), None\n",
    "        \n",
    "    # change view to adjust the gif orientation\n",
    "    viewup = [0, -8, 0] # pix3D_chairs\n",
    "    p = pv.Plotter()\n",
    "    trimesh_mesh = trimesh.Trimesh(mesh_vertices, faces=mesh_faces, process=False)\n",
    "    pv_mesh = pv.wrap(trimesh_mesh)\n",
    "    p.add_mesh(pv_mesh, scalars=mesh_colors, rgb=True, smooth_shading=True, cmap=\"rainbow\")\n",
    "    p.set_background(\"white\")\n",
    "    path = p.generate_orbital_path(factor=2.5, n_points=36, viewup=viewup, shift=0.2)\n",
    "    p.open_gif(\"{}.gif\".format(mesh_save_path))\n",
    "    p.orbit_on_path(path, write_frames=True, viewup=viewup, step=0.05)\n",
    "    p.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2693404d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using the following cell to visualize 3D meshes with deformationas as color.\n",
    "# need a display connect for open3d visualization\n",
    "# on open3d windown, press Ctrl + B to undo backface culling.\n",
    "\n",
    "# set output_dir to the data dump, and also provide the save directory.\n",
    "# path to output dump (generated by running evaluate.sh) eg: output/shapenet/reconstruction_stage_two/car/dump/\n",
    "output_dir = \"./\"\n",
    "gif_save_dir = './'\n",
    "\n",
    "input_image_files = sorted(glob.glob(os.path.join(output_dir, '*_image_input.png')))\n",
    "colored_mesh_ply_files = sorted(glob.glob(os.path.join(output_dir, '*_color_coded_mesh.ply')))\n",
    "non_colored_mesh_ply_files = sorted(glob.glob(os.path.join(output_dir, '*[0-9]_mesh.ply')))\n",
    "\n",
    "\n",
    "for idx in range(0, len(input_image_files)):\n",
    "    input_image = np.asarray(Image.open(input_image_files[idx]))\n",
    "    plt.imshow(input_image[:,:,:3])\n",
    "    plt.show()\n",
    "    \n",
    "    if len(colored_mesh_ply_files)>0:\n",
    "        visualize_mesh(colored_mesh_ply_files[idx], deformations_as_color=True)\n",
    "        render_mesh_gif(colored_mesh_ply_files[idx], mesh_save_path=\"{}/{}\".format(gif_save_dir, str(idx).zfill(6)), deformations_as_color=True)\n",
    "    elif len(non_colored_mesh_ply_files)>0:\n",
    "        visualize_mesh(non_colored_mesh_ply_files[idx])\n",
    "        render_mesh_gif(non_colored_mesh_ply_files[idx], mesh_save_path=\"{}/{}\".format(gif_save_dir, str(idx).zfill(6)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}