third_party/sam-3d-objects/notebook/demo_ss_encoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "originalKey": "10173bf8-2955-4905-abba-5d47338f1077"
   },
   "source": [
    "# SS Encoder/Decoder Round-Trip Demo\n",
    "\n",
    "This notebook loads GLB meshes, voxelizes them, runs them through the Sparse Structure VAE encoder and decoder, and visualizes the reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "cb6d7f66-1d98-40ec-9905-294c5a15a144",
    "output": {
     "id": 1609981380290449,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import trimesh\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "from sam3d_objects.model.backbone.tdfy_dit.models.sparse_structure_vae import (\n",
    "    SparseStructureEncoderTdfyWrapper,\n",
    "    SparseStructureDecoderTdfyWrapper,\n",
    ")\n",
    "\n",
    "# ============ CONFIG — edit these paths ============\n",
    "GLB_PATHS = [\n",
    "    \"YOUR GLB PATH\"\n",
    "]\n",
    "\n",
    "PROJECT_ROOT = \"./sam-3d-objects\"\n",
    "ENCODER_CKPT = os.path.join(PROJECT_ROOT, \"checkpoints/hf_weights/hf_weights/ss_encoder.ckpt\")\n",
    "DECODER_CKPT = os.path.join(PROJECT_ROOT, \"checkpoints/hf_weights/hf_weights/ss_decoder.ckpt\")\n",
    "OUTPUT_DIR = os.path.join(PROJECT_ROOT, \"notebook/output_ss_demo\")\n",
    "DEVICE = \"cuda:0\"\n",
    "RESOLUTION = 64\n",
    "\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "print(f\"Device: {DEVICE}\")\n",
    "print(f\"Resolution: {RESOLUTION}\")\n",
    "print(f\"Output dir: {OUTPUT_DIR}\")\n",
    "print(f\"GLB files: {len(GLB_PATHS)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "5adc274a-e54e-46ab-86c4-729b0246468f"
   },
   "outputs": [],
   "source": [
    "! which python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "6cb7b799-8906-48ef-b771-8103d8b0c954",
    "output": {
     "id": 1122478960948896,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "print(f\"torch {torch.__version__}, CUDA: {torch.cuda.is_available()}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "017aa939-b25d-42e8-b648-ee9edca5b12a"
   },
   "outputs": [],
   "source": [
    "def normalize_mesh_verts(verts):\n",
    "    \"\"\"Normalize mesh vertices to [-0.5, 0.5]^3.\"\"\"\n",
    "    vmin = verts.min(axis=0)\n",
    "    vmax = verts.max(axis=0)\n",
    "    center = (vmax + vmin) / 2.0\n",
    "    extent = vmax - vmin\n",
    "    max_extent = np.max(extent)\n",
    "    if max_extent == 0:\n",
    "        vertices = verts - center\n",
    "        scale = 1\n",
    "    else:\n",
    "        scale = 1.0 / max_extent\n",
    "        vertices = (verts - center) * scale\n",
    "    return vertices, scale, center\n",
    "\n",
    "\n",
    "def glb_to_voxels(glb_path, resolution=64, save_ply_path=None):\n",
    "    \"\"\"\n",
    "    Load a GLB file, voxelize it, and return an occupancy tensor.\n",
    "\n",
    "    Args:\n",
    "        glb_path: Path to the GLB file.\n",
    "        resolution: Voxel grid resolution (default 64).\n",
    "        save_ply_path: If set, save the normalized mesh as PLY.\n",
    "\n",
    "    Returns:\n",
    "        occupancy: Tensor of shape [1, resolution, resolution, resolution]\n",
    "        voxel_coords: Numpy array of voxel center coordinates\n",
    "    \"\"\"\n",
    "    # Load GLB\n",
    "    scene_or_mesh = trimesh.load(glb_path)\n",
    "    if isinstance(scene_or_mesh, trimesh.Scene):\n",
    "        mesh = scene_or_mesh.dump(concatenate=True)\n",
    "    else:\n",
    "        mesh = scene_or_mesh\n",
    "\n",
    "    verts = np.asarray(mesh.vertices)\n",
    "\n",
    "    # Y-up → Z-up rotation: (x, y, z) → (x, z, -y)\n",
    "    rot_matrix = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], dtype=np.float64)\n",
    "    verts = verts @ rot_matrix.T\n",
    "\n",
    "    # Normalize to [-0.5, 0.5]^3\n",
    "    verts, scale, center = normalize_mesh_verts(verts)\n",
    "\n",
    "    # Clamp to avoid boundary issues\n",
    "    verts = np.clip(verts, -0.5 + 1e-6, 0.5 - 1e-6)\n",
    "\n",
    "    # Update mesh vertices\n",
    "    mesh.vertices = verts\n",
    "\n",
    "    # Save normalized mesh if requested\n",
    "    if save_ply_path is not None:\n",
    "        mesh.export(save_ply_path)\n",
    "        print(f\"  Saved normalized mesh to {save_ply_path}\")\n",
    "\n",
    "    # Voxelize using trimesh\n",
    "    pitch = 1.0 / resolution\n",
    "    voxel_grid = mesh.voxelized(pitch)\n",
    "    # Fill interior voxels for watertight meshes, fall back to surface-only\n",
    "    try:\n",
    "        voxel_grid = voxel_grid.fill(method=\"holes\")\n",
    "    except Exception:\n",
    "        pass\n",
    "\n",
    "    # Extract voxel coordinates (centers of occupied voxels)\n",
    "    voxel_coords = voxel_grid.points\n",
    "\n",
    "    # Build occupancy tensor from voxel grid indices\n",
    "    sparse_indices = voxel_grid.sparse_indices  # [N, 3] integer grid indices\n",
    "    sparse_indices = np.clip(sparse_indices, 0, resolution - 1)\n",
    "    occupancy = torch.zeros(1, resolution, resolution, resolution, dtype=torch.float32)\n",
    "    occupancy[:, sparse_indices[:, 0], sparse_indices[:, 1], sparse_indices[:, 2]] = 1.0\n",
    "\n",
    "    return occupancy, voxel_coords\n",
    "\n",
    "\n",
    "print(\"Helper functions defined.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "71a7ec1d-c503-4190-b6de-12dc95439641"
   },
   "outputs": [],
   "source": [
    "# Load Encoder (config from ss_encoder.yaml)\n",
    "encoder = SparseStructureEncoderTdfyWrapper(\n",
    "    return_raw=True,\n",
    "    in_channels=1,\n",
    "    latent_channels=8,\n",
    "    channels=[32, 128, 512],\n",
    "    num_res_blocks=2,\n",
    "    num_res_blocks_middle=2,\n",
    "    pretrained_ckpt_path=ENCODER_CKPT,\n",
    ")\n",
    "encoder = encoder.to(DEVICE).eval()\n",
    "print(f\"Encoder loaded from {ENCODER_CKPT}\")\n",
    "print(f\"  Parameters: {sum(p.numel() for p in encoder.parameters()):,}\")\n",
    "\n",
    "# Load Decoder (config from ss_decoder.yaml)\n",
    "decoder = SparseStructureDecoderTdfyWrapper(\n",
    "    out_channels=1,\n",
    "    latent_channels=8,\n",
    "    channels=[512, 128, 32],\n",
    "    num_res_blocks=2,\n",
    "    num_res_blocks_middle=2,\n",
    "    reshape_input_to_cube=False,\n",
    "    pretrained_ckpt_path=DECODER_CKPT,\n",
    ")\n",
    "decoder = decoder.to(DEVICE).eval()\n",
    "print(f\"Decoder loaded from {DECODER_CKPT}\")\n",
    "print(f\"  Parameters: {sum(p.numel() for p in decoder.parameters()):,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "0c2e687d-178f-4809-80d3-28873a6c14e8"
   },
   "outputs": [],
   "source": [
    "# Run voxelization on all GLB files\n",
    "assert len(GLB_PATHS) > 0, \"Please add GLB file paths to GLB_PATHS in the config cell.\"\n",
    "\n",
    "voxel_results = {}  # name -> (occupancy, voxel_coords)\n",
    "\n",
    "for glb_path in GLB_PATHS:\n",
    "    name = Path(glb_path).stem\n",
    "    print(f\"\\nProcessing: {name}\")\n",
    "    ply_path = os.path.join(OUTPUT_DIR, f\"{name}_normalized.ply\")\n",
    "    occupancy, voxel_coords = glb_to_voxels(glb_path, resolution=RESOLUTION, save_ply_path=ply_path)\n",
    "    voxel_results[name] = (occupancy, voxel_coords)\n",
    "    num_voxels = int(occupancy.sum().item())\n",
    "    print(f\"  Occupancy shape: {occupancy.shape}\")\n",
    "    print(f\"  Occupied voxels: {num_voxels} / {RESOLUTION**3} ({100*num_voxels/RESOLUTION**3:.1f}%)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "067e884f-5e48-41fb-ae85-5e411eaa74f5"
   },
   "outputs": [],
   "source": [
    "# Encode → Decode round-trip\n",
    "recon_results = {}  # name -> (input_occ, recon_occ, latent_stats)\n",
    "\n",
    "for name, (occupancy, _) in voxel_results.items():\n",
    "    print(f\"\\n=== {name} ===\")\n",
    "    input_voxels = occupancy.unsqueeze(0).to(DEVICE)  # [1, 1, 64, 64, 64]\n",
    "    print(f\"Input shape: {input_voxels.shape}\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # Encode\n",
    "        enc_out = encoder(input_voxels)\n",
    "        z, mean, logvar = enc_out[\"z\"], enc_out[\"mean\"], enc_out[\"logvar\"]\n",
    "        print(f\"Latent z shape: {z.shape}\")\n",
    "        print(f\"Latent mean — mean: {mean.mean().item():.4f}, std: {mean.std().item():.4f}, \"\n",
    "              f\"min: {mean.min().item():.4f}, max: {mean.max().item():.4f}\")\n",
    "\n",
    "        # Decode from mean\n",
    "        logits = decoder(mean)\n",
    "        print(f\"Decoder output shape: {logits.shape}\")\n",
    "\n",
    "        # Threshold\n",
    "        recon_occ = (torch.sigmoid(logits) > 0.5).float()\n",
    "\n",
    "    # Compute IoU\n",
    "    input_bool = input_voxels.bool().cpu()\n",
    "    recon_bool = recon_occ.bool().cpu()\n",
    "    intersection = (input_bool & recon_bool).sum().item()\n",
    "    union = (input_bool | recon_bool).sum().item()\n",
    "    iou = intersection / union if union > 0 else 0.0\n",
    "    print(f\"IoU: {iou:.4f}  (intersection={intersection}, union={union})\")\n",
    "\n",
    "    recon_results[name] = (\n",
    "        occupancy,           # [1, 64, 64, 64]\n",
    "        recon_occ[0, 0].cpu(),  # [64, 64, 64]\n",
    "        {\"z_shape\": list(z.shape), \"iou\": iou},\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "827276d7-86e4-4ef4-ab3e-a8ec90171719"
   },
   "outputs": [],
   "source": [
    "# Visualization: side-by-side 3D scatter plots and PLY saving\n",
    "\n",
    "for name, (input_occ, recon_occ, stats) in recon_results.items():\n",
    "    # Get occupied voxel centers\n",
    "    input_coords = torch.nonzero(input_occ[0], as_tuple=False).numpy()  # [N, 3]\n",
    "    recon_coords = torch.nonzero(recon_occ, as_tuple=False).numpy()     # [M, 3]\n",
    "\n",
    "    # Convert grid indices to world coords: (idx + 0.5) / 64 - 0.5\n",
    "    input_pts = (input_coords + 0.5) / RESOLUTION - 0.5\n",
    "    recon_pts = (recon_coords + 0.5) / RESOLUTION - 0.5\n",
    "\n",
    "    # Save as PLY using trimesh\n",
    "    input_ply_path = os.path.join(OUTPUT_DIR, f\"{name}_input_voxels.ply\")\n",
    "    recon_ply_path = os.path.join(OUTPUT_DIR, f\"{name}_recon_voxels.ply\")\n",
    "\n",
    "    trimesh.PointCloud(input_pts).export(input_ply_path)\n",
    "    trimesh.PointCloud(recon_pts).export(recon_ply_path)\n",
    "\n",
    "    print(f\"Saved: {input_ply_path}, {recon_ply_path}\")\n",
    "\n",
    "    # Matplotlib 3D scatter\n",
    "    fig = plt.figure(figsize=(14, 6))\n",
    "\n",
    "    ax1 = fig.add_subplot(121, projection=\"3d\")\n",
    "    ax1.scatter(input_pts[:, 0], input_pts[:, 1], input_pts[:, 2], s=1, alpha=0.5)\n",
    "    ax1.set_title(f\"{name} — Input ({len(input_pts)} voxels)\")\n",
    "    ax1.set_xlim(-0.5, 0.5); ax1.set_ylim(-0.5, 0.5); ax1.set_zlim(-0.5, 0.5)\n",
    "\n",
    "    ax2 = fig.add_subplot(122, projection=\"3d\")\n",
    "    ax2.scatter(recon_pts[:, 0], recon_pts[:, 1], recon_pts[:, 2], s=1, alpha=0.5, c=\"orange\")\n",
    "    ax2.set_title(f\"{name} — Reconstructed ({len(recon_pts)} voxels, IoU={stats['iou']:.3f})\")\n",
    "    ax2.set_xlim(-0.5, 0.5); ax2.set_ylim(-0.5, 0.5); ax2.set_zlim(-0.5, 0.5)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "originalKey": "2f85e278-fecf-4cb6-bde3-c850b9464881"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "fileHeader": "",
  "fileUid": "35d6edac-96aa-45da-b73a-d89717b9c8f1",
  "isAdHoc": false,
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}