{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SAM 3D Body (3DB) Mesh Alignment to SAM 3D Object Scale\n",
    "\n",
    "This notebook processes a single 3DB mesh and aligns it to the SAM 3D Objects scale.\n",
    "\n",
    "**Input Data:**\n",
    "- `images/human_object/image.jpg` - Input image for MoGe\n",
    "- `meshes/human_object/3DB_results/mask_human.png` - Human mask\n",
    "- `meshes/human_object/3DB_results/human.ply` - Single 3DB mesh in OpenGL coordinates\n",
    "- `meshes/human_object/3DB_results/focal_length.json` - 3DB focal length\n",
    "\n",
    "**Output:**\n",
    "- `meshes/human_object/aligned_meshes/human_aligned.ply` - Aligned 3DB mesh in OpenGL coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from mesh_alignment import process_and_save_alignment\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f\"Using device: {device}\")\n",
    "PATH = os.getcwd()\n",
    "print(f\"Current working directory: {PATH}\")\n",
    "\n",
    "# Please inference the SAM 3D Body (3DB) Repo (https://github.com/facebookresearch/sam-3d-body) to get the 3DB Results\n",
    "image_path = f\"{PATH}/images/human_object/image.png\"\n",
    "mask_path = f\"{PATH}/meshes/human_object/3DB_results/mask_human.png\"\n",
    "mesh_path = f\"{PATH}/meshes/human_object/3DB_results/human.ply\"\n",
    "focal_length_json_path = f\"{PATH}/meshes/human_object/3DB_results/focal_length.json\"\n",
    "output_dir = f\"{PATH}/meshes/human_object/aligned_meshes\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Inference and Save SAM 3D Objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please inference SAM 3D Objects Repo with https://github.com/facebookresearch/sam-3d-objects/blob/main/notebook/demo_multi_object.ipynb\n",
    "# The above notebook will apply the generated layout to the generated objects, and same them as ply. \n",
    "# Then, this cell will load the posed SAM 3D Objects and transform them into the OpenGL coordinate system, which is the same system as SAM 3D Body. \n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "\n",
    "# Load PLY file\n",
    "input_path = 'gaussians/human_object_posed.ply'\n",
    "output_path = 'meshes/human_object/3Dfy_results/0.ply'\n",
    "mesh = o3d.io.read_point_cloud(input_path)\n",
    "points = np.asarray(mesh.points)\n",
    "\n",
    "# Transform to OpenGL coordinate system. \n",
    "points[:, [0, 2]] *= -1  # flip x and z\n",
    "mesh.points = o3d.utility.Vector3dVector(points)\n",
    "o3d.io.write_point_cloud(output_path, mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load and Display Input Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_image = Image.open(image_path)\n",
    "mask = Image.open(mask_path).convert('L')\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "axes[0].imshow(input_image)\n",
    "axes[0].set_title('Input Image')\n",
    "axes[0].axis('off')\n",
    "axes[1].imshow(mask, cmap='gray')\n",
    "axes[1].set_title('Mask')\n",
    "axes[1].axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Process and Save Aligned Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "success, output_mesh_path, result = process_and_save_alignment(\n",
    "    mesh_path=mesh_path,\n",
    "    mask_path=mask_path,\n",
    "    image_path=image_path,\n",
    "    output_dir=output_dir,\n",
    "    device=device,\n",
    "    focal_length_json_path=focal_length_json_path\n",
    ")\n",
    "\n",
    "if success:\n",
    "    print(f\"Alignment completed successfully! Output: {output_mesh_path}\")\n",
    "else:\n",
    "    print(\"Alignment failed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Interactive 3D Visualization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesh_alignment import visualize_meshes_interactive\n",
    "\n",
    "aligned_mesh_path = f\"{PATH}/meshes/human_object/aligned_meshes/human_aligned.ply\"\n",
    "dfy_mesh_path = f\"{PATH}/meshes/human_object/3Dfy_results/0.ply\"\n",
    "\n",
    "demo, combined_glb_path = visualize_meshes_interactive(\n",
    "    aligned_mesh_path=aligned_mesh_path,\n",
    "    dfy_mesh_path=dfy_mesh_path,\n",
    "    share=True\n",
    ")"
   ]
  }
 ],
 "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}