From 17dc3c529420c5f2f4795379d2da9890ceb69c28 Mon Sep 17 00:00:00 2001 From: Qiusheng Wu Date: Thu, 5 Feb 2026 22:56:07 -0500 Subject: [PATCH 1/2] feat: add SAM 3D Objects support for 3D reconstruction Adds a new module for reconstructing 3D objects from segmented masks using SAM 3D Objects (Meta's foundation model for 3D reconstruction from 2D images). Features: - Sam3DReconstructor class for single/multi-object reconstruction - Integration with SamGeo segmentation results via reconstruct_from_samgeo() - Gaussian splat output (PLY format) - Detailed installation instructions for the complex setup Requirements: - Linux 64-bit system - NVIDIA GPU with 32GB+ VRAM - HuggingFace authentication for model checkpoints Closes #461 --- docs/examples/sam3d.ipynb | 230 ++++++++++++++++++++++++++ docs/sam3d.md | 3 + mkdocs.yml | 2 + samgeo/__init__.py | 7 + samgeo/sam3d.py | 334 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 576 insertions(+) create mode 100644 docs/examples/sam3d.ipynb create mode 100644 docs/sam3d.md create mode 100644 samgeo/sam3d.py diff --git a/docs/examples/sam3d.ipynb b/docs/examples/sam3d.ipynb new file mode 100644 index 00000000..443fa0ce --- /dev/null +++ b/docs/examples/sam3d.ipynb @@ -0,0 +1,230 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 3D Object Reconstruction with SAM 3D\n", + "\n", + "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/segment-geospatial/blob/main/docs/examples/sam3d.ipynb)\n", + "\n", + "This notebook demonstrates how to use SAM 3D Objects to reconstruct 3D models from segmented objects in images.\n", + "\n", + "## Overview\n", + "\n", + "[SAM 3D Objects](https://github.com/facebookresearch/sam-3d-objects) is a foundation model from Meta that converts masked objects in images into 3D models with pose, shape, texture, and layout. It excels in real-world scenarios with occlusion and clutter.\n", + "\n", + "**Requirements:**\n", + "- Linux 64-bit system\n", + "- NVIDIA GPU with at least 32GB VRAM\n", + "- HuggingFace authentication for checkpoint access\n", + "\n", + "**Reference:**\n", + "SAM 3D Team (2025). SAM 3D: 3Dfy Anything in Images. https://arxiv.org/abs/2511.16624" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Installation\n", + "\n", + "SAM 3D Objects requires a separate installation. Print the instructions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from samgeo.sam3d import print_install_instructions\n", + "\n", + "print_install_instructions()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup Environment Variable\n", + "\n", + "After installing SAM 3D Objects, set the path to the repository:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Set the path to your sam-3d-objects installation\n", + "# os.environ[\"SAM3D_PATH\"] = \"/path/to/sam-3d-objects\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Basic Usage\n", + "\n", + "Reconstruct a 3D object from an image and mask:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from samgeo.sam3d import Sam3DReconstructor\n", + "\n", + "# Initialize the reconstructor\n", + "# reconstructor = Sam3DReconstructor()\n", + "\n", + "# Reconstruct a single object\n", + "# output = reconstructor.reconstruct(\n", + "# image=\"path/to/image.png\",\n", + "# mask=\"path/to/mask.png\",\n", + "# seed=42,\n", + "# )\n", + "\n", + "# Save the Gaussian splat\n", + "# reconstructor.save_ply(output, \"object.ply\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Workflow: Segment with SAM, Reconstruct with SAM 3D\n", + "\n", + "A typical workflow combines SAM/SAM2/SAM3 for segmentation with SAM 3D for 3D reconstruction:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Step 1: Segment objects using SamGeo\n", + "# from samgeo import SamGeo\n", + "# \n", + "# sam = SamGeo()\n", + "# sam.set_image(\"image.tif\")\n", + "# sam.generate(output=\"masks.gpkg\")\n", + "\n", + "# Step 2: Reconstruct 3D models from segmented objects\n", + "# from samgeo.sam3d import reconstruct_from_samgeo\n", + "# import geopandas as gpd\n", + "# \n", + "# masks = gpd.read_file(\"masks.gpkg\")\n", + "# ply_files = reconstruct_from_samgeo(\n", + "# samgeo_result=masks,\n", + "# image_path=\"image.tif\",\n", + "# output_dir=\"./3d_models\",\n", + "# max_objects=5, # Limit number of objects\n", + "# )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multiple Object Reconstruction\n", + "\n", + "Reconstruct multiple objects from a single image:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# reconstructor = Sam3DReconstructor()\n", + "\n", + "# masks = [\"mask1.png\", \"mask2.png\", \"mask3.png\"]\n", + "# outputs = reconstructor.reconstruct_multiple(\n", + "# image=\"image.png\",\n", + "# masks=masks,\n", + "# )\n", + "\n", + "# for i, output in enumerate(outputs):\n", + "# reconstructor.save_ply(output, f\"object_{i}.ply\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Output Formats\n", + "\n", + "SAM 3D Objects outputs Gaussian splats which can be:\n", + "- Saved as PLY files for viewing in 3D viewers\n", + "- Converted to meshes for use in 3D software\n", + "- Visualized in web-based viewers\n", + "\n", + "For viewing PLY files, you can use:\n", + "- [MeshLab](https://www.meshlab.net/)\n", + "- [CloudCompare](https://www.cloudcompare.org/)\n", + "- [Three.js](https://threejs.org/) web viewer" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Tips and Best Practices\n", + "\n", + "1. **GPU Memory**: SAM 3D requires at least 32GB VRAM. Use smaller images or process objects one at a time if you have memory constraints.\n", + "\n", + "2. **Mask Quality**: Better segmentation masks lead to better 3D reconstructions. Use SAM2 or SAM3 for high-quality masks.\n", + "\n", + "3. **Object Selection**: SAM 3D works best on:\n", + " - Complete, unoccluded objects\n", + " - Objects with clear boundaries\n", + " - Objects with sufficient texture\n", + "\n", + "4. **Seed Consistency**: Use the same seed for reproducible results.\n", + "\n", + "5. **Batch Processing**: For multiple objects, use `reconstruct_multiple()` or `reconstruct_from_samgeo()` for efficiency." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "- [SAM 3D Objects GitHub](https://github.com/facebookresearch/sam-3d-objects)\n", + "- [SAM 3D Website](https://ai.meta.com/sam3d/)\n", + "- [SAM 3D Paper](https://arxiv.org/abs/2511.16624)\n", + "- [SAM 3D Demo](https://www.aidemos.meta.com/segment-anything/editor/convert-image-to-3d)\n", + "- [SAM 3D Body](https://github.com/facebookresearch/sam-3d-body) (for human mesh recovery)" + ] + } + ], + "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", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/sam3d.md b/docs/sam3d.md new file mode 100644 index 00000000..79f85d70 --- /dev/null +++ b/docs/sam3d.md @@ -0,0 +1,3 @@ +# sam3d module + +::: samgeo.sam3d diff --git a/mkdocs.yml b/mkdocs.yml index 58206728..8ef54e95 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -83,6 +83,7 @@ nav: - examples/sam3_box_prompts.ipynb - examples/sam3_tiled_segmentation.ipynb - examples/detectree2.ipynb + - examples/sam3d.ipynb - Workshops: - workshops/purdue.ipynb - workshops/cn_workshop.ipynb @@ -98,4 +99,5 @@ nav: - hq_sam module: hq_sam.md - text_sam module: text_sam.md - detectree2 module: detectree2.md + - sam3d module: sam3d.md # - fer module: fer.md diff --git a/samgeo/__init__.py b/samgeo/__init__.py index 308f1364..4cbe4608 100644 --- a/samgeo/__init__.py +++ b/samgeo/__init__.py @@ -21,3 +21,10 @@ ) except ImportError: pass # detectree2 not installed + +# SAM 3D Objects support for 3D reconstruction +from .sam3d import ( + Sam3DReconstructor, + print_install_instructions as sam3d_install_instructions, + reconstruct_from_samgeo, +) diff --git a/samgeo/sam3d.py b/samgeo/sam3d.py new file mode 100644 index 00000000..3f5408e9 --- /dev/null +++ b/samgeo/sam3d.py @@ -0,0 +1,334 @@ +"""SAM 3D Objects module for 3D reconstruction from segmented masks. + +This module provides tools for reconstructing 3D objects from 2D images using +SAM 3D Objects, a foundation model from Meta that converts masked objects in +images into 3D models with pose, shape, texture, and layout. + +Note: + SAM 3D Objects requires: + - Linux 64-bit system + - NVIDIA GPU with at least 32GB VRAM + - HuggingFace authentication for checkpoint access + +Reference: + SAM 3D Team (2025). SAM 3D: 3Dfy Anything in Images. + https://arxiv.org/abs/2511.16624 + + Repository: https://github.com/facebookresearch/sam-3d-objects + Website: https://ai.meta.com/sam3d/ +""" + +import logging +import os +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np + +logger = logging.getLogger(__name__) + +# Installation instructions +SAM3D_INSTALL_INSTRUCTIONS = """ +SAM 3D Objects Installation Instructions +======================================== + +Prerequisites: +- Linux 64-bit system +- NVIDIA GPU with at least 32GB VRAM +- Conda or Mamba package manager + +Step 1: Clone the repository +---------------------------- +git clone https://github.com/facebookresearch/sam-3d-objects.git +cd sam-3d-objects + +Step 2: Create environment +-------------------------- +mamba env create -f environments/default.yml +mamba activate sam3d-objects + +Step 3: Install dependencies +---------------------------- +export PIP_EXTRA_INDEX_URL="https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu121" +pip install -e '.[dev]' +pip install -e '.[p3d]' +export PIP_FIND_LINKS="https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.5.1_cu121.html" +pip install -e '.[inference]' +./patching/hydra + +Step 4: Download checkpoints (requires HuggingFace authentication) +----------------------------------------------------------------- +# First, request access at: https://huggingface.co/facebook/sam-3d-objects +# Then authenticate: huggingface-cli login + +pip install 'huggingface-hub[cli]<1.0' +TAG=hf +huggingface-cli download \\ + --repo-type model \\ + --local-dir checkpoints/${TAG}-download \\ + --max-workers 1 \\ + facebook/sam-3d-objects +mv checkpoints/${TAG}-download/checkpoints checkpoints/${TAG} +rm -rf checkpoints/${TAG}-download + +For more details, see: https://github.com/facebookresearch/sam-3d-objects/blob/main/doc/setup.md +""" + + +def _check_sam3d(): + """Check if SAM 3D Objects is installed and raise informative error if not.""" + try: + # Try to import from the sam-3d-objects package + import sys + + # Check if the sam-3d-objects notebook inference module is available + # The package doesn't have a standard import, so we check for the inference module + sam3d_path = os.environ.get("SAM3D_PATH") + if sam3d_path: + sys.path.insert(0, os.path.join(sam3d_path, "notebook")) + + from inference import Inference + + return Inference + except ImportError: + raise ImportError( + "SAM 3D Objects is not installed or not configured properly.\n\n" + f"{SAM3D_INSTALL_INSTRUCTIONS}\n\n" + "After installation, set the SAM3D_PATH environment variable:\n" + " export SAM3D_PATH=/path/to/sam-3d-objects\n" + ) + + +def print_install_instructions(): + """Print SAM 3D Objects installation instructions.""" + print(SAM3D_INSTALL_INSTRUCTIONS) + + +class Sam3DReconstructor: + """Class for 3D object reconstruction using SAM 3D Objects. + + This class provides methods for converting segmented masks into 3D models + (Gaussian splats or meshes) using the SAM 3D Objects model. + + Attributes: + inference: The SAM 3D inference object. + config_path: Path to the pipeline configuration. + + Example: + >>> from samgeo.sam3d import Sam3DReconstructor + >>> reconstructor = Sam3DReconstructor() + >>> output = reconstructor.reconstruct("image.png", "mask.png") + >>> output["gs"].save_ply("object.ply") + """ + + def __init__( + self, + sam3d_path: Optional[str] = None, + config_tag: str = "hf", + compile_model: bool = False, + ) -> None: + """Initialize the Sam3DReconstructor. + + Args: + sam3d_path: Path to the sam-3d-objects repository. If None, uses + the SAM3D_PATH environment variable. + config_tag: Configuration tag for the model ('hf' for HuggingFace). + compile_model: Whether to compile the model for faster inference. + """ + self.sam3d_path = sam3d_path or os.environ.get("SAM3D_PATH") + if not self.sam3d_path: + raise ValueError( + "SAM3D_PATH not set. Please provide sam3d_path or set the " + "SAM3D_PATH environment variable to the sam-3d-objects directory." + ) + + self.config_tag = config_tag + self.compile_model = compile_model + self._inference = None + + # Add the notebook path to sys.path for imports + import sys + + notebook_path = os.path.join(self.sam3d_path, "notebook") + if notebook_path not in sys.path: + sys.path.insert(0, notebook_path) + + logger.info(f"Sam3DReconstructor initialized with path: {self.sam3d_path}") + + def _setup_inference(self) -> None: + """Set up the SAM 3D inference pipeline.""" + if self._inference is not None: + return + + from inference import Inference + + config_path = os.path.join( + self.sam3d_path, "checkpoints", self.config_tag, "pipeline.yaml" + ) + + if not os.path.exists(config_path): + raise FileNotFoundError( + f"Config file not found: {config_path}\n" + "Please download the checkpoints first. See print_install_instructions()." + ) + + logger.info(f"Loading SAM 3D model from: {config_path}") + self._inference = Inference(config_path, compile=self.compile_model) + logger.info("SAM 3D model loaded successfully") + + def reconstruct( + self, + image: Union[str, np.ndarray, "PIL.Image.Image"], + mask: Union[str, np.ndarray, "PIL.Image.Image"], + seed: int = 42, + **kwargs: Any, + ) -> Dict[str, Any]: + """Reconstruct a 3D object from an image and mask. + + Args: + image: Input image (path, numpy array, or PIL Image). + mask: Binary mask for the object (path, numpy array, or PIL Image). + seed: Random seed for reproducibility. + **kwargs: Additional arguments passed to the inference. + + Returns: + Dictionary containing: + - 'gs': Gaussian splat object (can be saved with .save_ply()) + - Additional outputs from the model + """ + self._setup_inference() + + from inference import load_image, load_single_mask + + # Load image if path + if isinstance(image, str): + image = load_image(image) + + # Load mask if path + if isinstance(mask, str): + mask = load_single_mask(os.path.dirname(mask), index=0) + elif isinstance(mask, np.ndarray): + # Convert numpy mask to the expected format + from PIL import Image + + if mask.dtype == bool: + mask = mask.astype(np.uint8) * 255 + mask = Image.fromarray(mask) + + logger.info("Running 3D reconstruction...") + output = self._inference(image, mask, seed=seed, **kwargs) + logger.info("3D reconstruction complete") + + return output + + def reconstruct_multiple( + self, + image: Union[str, np.ndarray, "PIL.Image.Image"], + masks: List[Union[str, np.ndarray, "PIL.Image.Image"]], + seed: int = 42, + **kwargs: Any, + ) -> List[Dict[str, Any]]: + """Reconstruct multiple 3D objects from an image and multiple masks. + + Args: + image: Input image (path, numpy array, or PIL Image). + masks: List of binary masks for objects. + seed: Random seed for reproducibility. + **kwargs: Additional arguments passed to the inference. + + Returns: + List of dictionaries, each containing reconstruction outputs. + """ + results = [] + for i, mask in enumerate(masks): + logger.info(f"Reconstructing object {i + 1}/{len(masks)}...") + output = self.reconstruct(image, mask, seed=seed + i, **kwargs) + results.append(output) + return results + + def save_ply( + self, + output: Dict[str, Any], + output_path: str, + ) -> str: + """Save the Gaussian splat to a PLY file. + + Args: + output: Output dictionary from reconstruct(). + output_path: Path to save the PLY file. + + Returns: + Path to the saved PLY file. + """ + if "gs" not in output: + raise ValueError("Output does not contain Gaussian splat ('gs')") + + output["gs"].save_ply(output_path) + logger.info(f"Gaussian splat saved to: {output_path}") + return output_path + + +def reconstruct_from_samgeo( + samgeo_result: "gpd.GeoDataFrame", + image_path: str, + output_dir: str, + sam3d_path: Optional[str] = None, + max_objects: int = 10, + seed: int = 42, +) -> List[str]: + """Reconstruct 3D objects from SamGeo segmentation results. + + This function takes segmentation results from SamGeo (or similar) and + reconstructs 3D models for each segmented object. + + Args: + samgeo_result: GeoDataFrame with segmentation masks. + image_path: Path to the original image. + output_dir: Directory to save the 3D models. + sam3d_path: Path to sam-3d-objects. Uses SAM3D_PATH env var if None. + max_objects: Maximum number of objects to reconstruct. + seed: Random seed for reproducibility. + + Returns: + List of paths to the saved PLY files. + """ + import geopandas as gpd + import rasterio + from rasterio.features import rasterize + + os.makedirs(output_dir, exist_ok=True) + + # Initialize reconstructor + reconstructor = Sam3DReconstructor(sam3d_path=sam3d_path) + + # Read image bounds for mask rasterization + with rasterio.open(image_path) as src: + transform = src.transform + shape = (src.height, src.width) + + output_files = [] + num_objects = min(len(samgeo_result), max_objects) + + for i in range(num_objects): + geom = samgeo_result.iloc[i].geometry + # Rasterize the geometry to create a mask + mask = rasterize( + [(geom, 1)], + out_shape=shape, + transform=transform, + fill=0, + dtype=np.uint8, + ) + mask = (mask > 0).astype(np.uint8) * 255 + + try: + output = reconstructor.reconstruct(image_path, mask, seed=seed + i) + output_path = os.path.join(output_dir, f"object_{i:03d}.ply") + reconstructor.save_ply(output, output_path) + output_files.append(output_path) + except Exception as e: + logger.warning(f"Failed to reconstruct object {i}: {e}") + continue + + logger.info(f"Reconstructed {len(output_files)}/{num_objects} objects") + return output_files From d5918ef0a4026dbfeeb6ab5c03c61f2e9b5af47e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Feb 2026 03:57:57 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/examples/sam3d.ipynb | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/examples/sam3d.ipynb b/docs/examples/sam3d.ipynb index 443fa0ce..0ba192bb 100644 --- a/docs/examples/sam3d.ipynb +++ b/docs/examples/sam3d.ipynb @@ -58,7 +58,6 @@ "metadata": {}, "outputs": [], "source": [ - "import os\n", "\n", "# Set the path to your sam-3d-objects installation\n", "# os.environ[\"SAM3D_PATH\"] = \"/path/to/sam-3d-objects\"" @@ -79,7 +78,6 @@ "metadata": {}, "outputs": [], "source": [ - "from samgeo.sam3d import Sam3DReconstructor\n", "\n", "# Initialize the reconstructor\n", "# reconstructor = Sam3DReconstructor()\n",