dev-bjoern/sam3d-objects-mcp

1"""Minimal SAM 3D Objects inference wrapper.2 3Replaces upstream notebook/inference.py, whose module-level imports pull in4kaolin.visualize, SceneVisualizer and plotly — none of which are needed to5run the pipeline. Import this module only when sam-3d-objects is on sys.path6and a GPU context is available.7"""8import os9from typing import Optional, Union10 11import numpy as np12from PIL import Image13from omegaconf import OmegaConf14from hydra.utils import instantiate15 16import sam3d_objects  # noqa: F401  guarded by LIDRA_SKIP_INIT17 18# The attention modules read ATTN_BACKEND/SPARSE_ATTN_BACKEND from the19# environment exactly once, at import time. inference_pipeline's20# set_attention_backend() flips the env to flash_attn on datacenter GPUs,21# so the modules must be imported BEFORE the pipeline to stay on sdpa.22import sam3d_objects.model.backbone.tdfy_dit.modules.attention  # noqa: F40123import sam3d_objects.model.backbone.tdfy_dit.modules.sparse  # noqa: F40124 25from sam3d_objects.pipeline.inference_pipeline_pointmap import InferencePipelinePointMap26 27 28class SAM3DInference:29    def __init__(self, config_file: str, compile: bool = False):30        config = OmegaConf.load(config_file)31        config.rendering_engine = "pytorch3d"  # disable nvdiffrast32        config.compile_model = compile33        config.workspace_dir = os.path.dirname(config_file)34        self._pipeline: InferencePipelinePointMap = instantiate(config)35 36    @staticmethod37    def merge_mask_to_rgba(image: np.ndarray, mask: np.ndarray) -> np.ndarray:38        mask = mask.astype(np.uint8) * 25539        return np.concatenate([image[..., :3], mask[..., None]], axis=-1)40 41    def __call__(42        self,43        image: Union[Image.Image, np.ndarray],44        mask: Optional[Union[Image.Image, np.ndarray]],45        seed: Optional[int] = None,46        pointmap=None,47    ) -> dict:48        image = self.merge_mask_to_rgba(np.asarray(image), np.asarray(mask))49        return self._pipeline.run(50            image,51            None,52            seed,53            stage1_only=False,54            with_mesh_postprocess=False,55            with_texture_baking=False,56            with_layout_postprocess=False,57            use_vertex_color=True,58            stage1_inference_steps=None,59            pointmap=pointmap,60        )61