{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Imports and Model Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "parent_dir = os.path.dirname(os.getcwd()) \n",
    "sys.path.insert(0, parent_dir)\n",
    "\n",
    "from utils import (\n",
    "    setup_sam_3d_body, setup_visualizer, \n",
    "    visualize_2d_results, visualize_3d_mesh, save_mesh_results, \n",
    "    display_results_grid, process_image_with_mask\n",
    ")\n",
    "\n",
    "# Set up SAM 3D Body estimator\n",
    "estimator = setup_sam_3d_body(hf_repo_id=\"facebook/sam-3d-body-dinov3\")\n",
    "# Set up visualizer\n",
    "visualizer = setup_visualizer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Process Image and Get Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Load and process the image\n",
    "image_path = \"images/dancing.jpg\"  # Relative to notebook folder\n",
    "img_cv2 = cv2.imread(image_path)\n",
    "img_rgb = cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# Process the image with SAM 3D Body\n",
    "print(\"Processing image with SAM 3D Body...\")\n",
    "outputs = estimator.process_one_image(image_path)\n",
    "\n",
    "print(f\"Number of people detected: {len(outputs)}\")\n",
    "print(f\"Output keys for first person: {list(outputs[0].keys()) if outputs else 'No people detected'}\")\n",
    "\n",
    "# Display the original image\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.imshow(img_rgb)\n",
    "plt.axis('off')\n",
    "plt.title('Original Image')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 2D Visualization - Keypoints and Bounding Boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize 2D results using utils\n",
    "if outputs:\n",
    "    vis_results = visualize_2d_results(img_cv2, outputs, visualizer)\n",
    "    \n",
    "    # Display results using grid function\n",
    "    titles = [f'Person {i} - 2D Keypoints & BBox' for i in range(len(vis_results))]\n",
    "    display_results_grid(vis_results, titles, figsize_per_image=(6, 6))\n",
    "else:\n",
    "    print(\"No people detected in the image\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 3D Mesh Visualization - Overlay and Side View"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if outputs:\n",
    "    mesh_results = visualize_3d_mesh(img_cv2, outputs, estimator.faces)\n",
    "    \n",
    "    # Display results\n",
    "    for i, combined_img in enumerate(mesh_results):\n",
    "        combined_rgb = cv2.cvtColor(combined_img, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        plt.figure(figsize=(20, 5))\n",
    "        plt.imshow(combined_rgb)\n",
    "        plt.title(f'Person {i}: Original | Mesh Overlay | Front View | Side View')\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "else:\n",
    "    print(\"No people detected for 3D mesh visualization\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Save 3D Mesh Files and Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if outputs:\n",
    "    # Get image name without extension\n",
    "    image_name = os.path.splitext(os.path.basename(image_path))[0]\n",
    "    \n",
    "    # Create output directory\n",
    "    output_dir = f\"output/{image_name}\"\n",
    "\n",
    "    # Save all results (PLY meshes, overlay images, bbox images)\n",
    "    ply_files = save_mesh_results(img_cv2, outputs, estimator.faces, output_dir, image_name)\n",
    "    \n",
    "    print(f\"\\n=== Saved Results for {image_name} ===\")\n",
    "    print(f\"Output directory: {output_dir}\")\n",
    "    print(f\"Number of PLY files created: {len(ply_files)}\")\n",
    "    \n",
    "else:\n",
    "    print(\"No results to save - no people detected\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Mask-Based Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load mask and run inference\n",
    "mask_path = \"images/dancing_mask.png\"\n",
    "\n",
    "if os.path.exists(mask_path):\n",
    "    # Load and display the mask\n",
    "    mask_img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.imshow(mask_img, cmap='gray')\n",
    "    plt.title('External Mask')\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    # Process with external mask\n",
    "    mask_outputs = process_image_with_mask(estimator, image_path, mask_path)\n",
    "    \n",
    "    # Visualize and save results\n",
    "    if mask_outputs:\n",
    "        mask_mesh_results = visualize_3d_mesh(img_cv2, mask_outputs, estimator.faces)\n",
    "        \n",
    "        for i, combined_img in enumerate(mask_mesh_results):\n",
    "            combined_rgb = cv2.cvtColor(combined_img, cv2.COLOR_BGR2RGB)\n",
    "            plt.figure(figsize=(20, 5))\n",
    "            plt.imshow(combined_rgb)\n",
    "            plt.title(f'Mask-Based Person {i}: Original | Mesh Overlay | Front View | Side View')\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "            \n",
    "        # Save results\n",
    "        mask_output_dir = f\"output/mask_based_{image_name}\"\n",
    "        mask_ply_files = save_mesh_results(img_cv2, mask_outputs, estimator.faces, mask_output_dir, f\"mask_{image_name}\")\n",
    "        print(f\"Saved mask-based results to: {mask_output_dir}\")\n",
    "    else:\n",
    "        print(\"No people detected with mask-based approach\")\n",
    "        \n",
    "else:\n",
    "    print(f\"Mask file not found: {mask_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3po_push",
   "language": "python",
   "name": "3po_push"
  },
  "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}