visualization_scripts/texture_transfer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "98eb44f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob, math, random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from scipy.spatial import KDTree\n",
    "import matplotlib.pyplot as plt\n",
    "import trimesh\n",
    "import open3d as o3d\n",
    "from open3d_utils import make_pcd, make_line_set, make_mesh\n",
    "# export DISPLAY for remote visualization\n",
    "# os.environ['DISPLAY'] = \":99\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "85b74ad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transfer_texture(output_dir, target_object_idx, source_mesh_data):\n",
    "    # load target mesh\n",
    "    target_mesh = o3d.io.read_triangle_mesh(os.path.join(output_dir, '{}_mesh.ply'.format(target_object_idx)))        \n",
    "    target_vertices = np.asarray(target_mesh.vertices)\n",
    "    target_faces = np.asarray(target_mesh.triangles)\n",
    "    target_reconstructed_pointcloud = np.load(os.path.join(output_dir, '{}_point_cloud_pred.npy'.format(target_object_idx)))\n",
    "    target_reconstructed_canonical_pointcloud = np.load(os.path.join(output_dir, '{}_point_cloud_pred_canonical.npy'.format(target_object_idx)))\n",
    "\n",
    "    \n",
    "    # transfer color from source mesh vertices --> source mesh 3D point cloud and source mesh canonical 3D point cloud\n",
    "    source_vertices_kdtree = KDTree(source_mesh_data[\"source_vertices\"])\n",
    "    source_reconstructed_pointcloud_mapping_idx = source_vertices_kdtree.query(source_mesh_data[\"source_reconstructed_pointcloud\"], k=1)[1]\n",
    "    source_reconstructed_pointcloud_color = source_mesh_data[\"source_colors\"][source_reconstructed_pointcloud_mapping_idx]\n",
    "    \n",
    "    # map target mesh vertices --> corresponding canonical 3D point cloud.\n",
    "    target_reconstructed_pointcloud_kdtree = KDTree(target_reconstructed_pointcloud)\n",
    "    target_vertices_mapping_idx = target_reconstructed_pointcloud_kdtree.query(target_vertices, k=1)[1]\n",
    "    target_vertices_canonical_point = target_reconstructed_canonical_pointcloud[target_vertices_mapping_idx]\n",
    "      \n",
    "    # transfer color from source point cloud to target point cloud\n",
    "    source_canonical_kdtree = KDTree(source_mesh_data[\"source_reconstructed_canonical_pointcloud\"])\n",
    "    target_vertices_canonical_mapping_idx = source_canonical_kdtree.query(target_vertices_canonical_point, k=1)[1]\n",
    "    target_vertices_color = source_reconstructed_pointcloud_color[target_vertices_canonical_mapping_idx]\n",
    "    \n",
    "    # visualize target input image\n",
    "    plt.imshow(np.asarray(Image.open(os.path.join(output_dir, '{}_image_input.png'.format(target_object_idx)))))\n",
    "    plt.show()\n",
    "    # visualize colored source mesh\n",
    "    o3d.visualization.draw_geometries([make_mesh(source_mesh_data[\"source_vertices\"], source_mesh_data[\"source_faces\"], source_mesh_data[\"source_colors\"][:,:3])])\n",
    "    # visualize uncolored target mesh\n",
    "    o3d.visualization.draw_geometries([make_mesh(target_vertices, target_faces)])\n",
    "    # visualize colored target mesh\n",
    "    o3d.visualization.draw_geometries([make_mesh(target_vertices, target_faces,target_vertices_color[:,:3])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bec8ecd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# two steps to trasnfer texture: set the `output_dir` and set `texture_mesh_dir`\n",
    "\n",
    "def load_source_mesh(textured_mesh_dir):\n",
    "    # load source textured mesh (source mesh is first reconstructed by TARS, then painted using meshlab)\n",
    "    source_mesh = o3d.io.read_triangle_mesh(os.path.join(textured_mesh_dir, 'textured_mesh.ply'))\n",
    "    source_vertices = np.asarray(source_mesh.vertices)\n",
    "    source_faces = np.asarray(source_mesh.triangles)\n",
    "    source_colors = np.asarray(source_mesh.vertex_colors)\n",
    "    source_colors = source_colors / 1.2\n",
    "    source_reconstructed_pointcloud = np.load(os.path.join(textured_mesh_dir, 'point_cloud_pred.npy'))\n",
    "    source_reconstructed_canonical_pointcloud = np.load(os.path.join(textured_mesh_dir, 'point_cloud_pred_canonical.npy'))\n",
    "    source_mesh_data = {}\n",
    "    source_mesh_data[\"source_vertices\"] = source_vertices\n",
    "    source_mesh_data[\"source_faces\"] = source_faces\n",
    "    source_mesh_data[\"source_colors\"] = source_colors\n",
    "    source_mesh_data[\"source_reconstructed_pointcloud\"] = source_reconstructed_pointcloud\n",
    "    source_mesh_data[\"source_reconstructed_canonical_pointcloud\"] = source_reconstructed_canonical_pointcloud\n",
    "    return source_mesh_data\n",
    "\n",
    "\n",
    "def main():\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/pascal3d/reconstruction_stage_two/chair/dump/\n",
    "    output_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",
    "    # Please follow the instructions in \"\" to texture new source meshes yourself and then transfer to other meshes.\n",
    "    # To use the one of the shared textured mesh as the source mesh, run the following:\n",
    "    # texture_mesh_dir = <path_to_dumped_textured_meshes>\n",
    "    # texture_mesh_dir = '../textured_meshes/70_mesh/'\n",
    "    # texture_mesh_dir = '../textured_meshes/288_mesh/'\n",
    "    texture_mesh_dir = '../textured_meshes/448_mesh/'\n",
    "    source_mesh_data = load_source_mesh(texture_mesh_dir)\n",
    "\n",
    "    # to transfer source mesh texture to all other meshes\n",
    "    for img_path in input_image_files:\n",
    "        target_idx = int(img_path.split('/')[-1].split('_')[0])\n",
    "        transfer_texture(output_dir, target_idx, source_mesh_data)\n",
    "        \n",
    "\n",
    "main()"
   ]
  }
 ],
 "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
}