Align the generated 3D scene's rendering with the input image

[snowflakewang]

Hello, thank you for your brilliant work!

I have input an image to the SAM 3D model and get the generated 3D scene. How can I render the 3D scene to get a rendered image that is aligned with the input image? i.e. How can I get the estimated intrinsics of the input image?

Thank you!

[RafailFridman]

Hi
I just created a script that seems to do that
I took parts from this pr and very importantly I set "with_layout_postprocess=True" in https://github.com/facebookresearch/sam-3d-objects/blob/main/notebook/inference.py#L116

It saves the transformed obj file together with the render overlay to see how well it is aligned to an image

If you will find some issues with the script or some ways to improve it, please tell

import argparse
from pathlib import Path
import copy
import torch
import numpy as np
from PIL import Image
from pytorch3d.transforms import quaternion_to_matrix, Transform3d
import open3d as o3d


from inference import Inference, load_image
from moge.model.v1 import MoGeModel

R_yup_to_zup = torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]], dtype=torch.float32)

R_flip_z = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -1]], dtype=torch.float32)

R_pytorch3d_to_zup = torch.tensor(
    [
        [-1, 0, 0],  # Keep X flipped (Matches your original handedness)
        [0, 0, 1],  # Map internal Y (Up) to Output Z (Up)
        [0, 1, 0],  # Map internal Z (Forward) to Output Y (Forward)
    ],
    dtype=torch.float32,
)


def get_moge_pointcloud(image_tensor, device="cuda"):
    """Generate MoGe point cloud from image tensor."""
    moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(device)
    moge_model.eval()
    with torch.no_grad():
        moge_output = moge_model.infer(image_tensor)
    return moge_output


def denormalize_intrinsics(norm_K, height, width):
    """Convert normalized intrinsics to absolute pixel coordinates."""
    cx_norm, cy_norm = norm_K[0, 2], norm_K[1, 2]
    fx_norm, fy_norm = norm_K[0, 0], norm_K[1, 1]

    fx_abs = fx_norm * width
    fy_abs = fy_norm * height
    cx_abs = cx_norm * width
    cy_abs = cy_norm * height
    fx_abs = fy_abs

    return np.array([[fx_abs, 0.0, cx_abs], [0.0, fy_abs, cy_abs], [0.0, 0.0, 1.0]])


def render_mesh_with_intrinsics(obj_path, intrinsics, width, height):

    renderer = o3d.visualization.rendering.OffscreenRenderer(width, height)

    renderer.scene.set_background([0.0, 0.0, 0.0, 1.0])
    renderer.scene.set_lighting(
        renderer.scene.LightingProfile.NO_SHADOWS, (0.577, -0.577, -0.577)
    )

    mesh = o3d.io.read_triangle_mesh(obj_path)
    if not mesh.has_vertex_normals():
        mesh.compute_vertex_normals()

    mat = o3d.visualization.rendering.MaterialRecord()
    mat.shader = "defaultLit"
    renderer.scene.add_geometry("mesh", mesh, mat)

    near_plane = 0.01
    far_plane = 100.0

    K = np.array(intrinsics, dtype=np.float64)

    renderer.scene.camera.set_projection(K, near_plane, far_plane, width, height)

    eye = [0, 0, 0]
    center = [0, 1, 0]
    up = [0, 0, 1]

    renderer.scene.camera.look_at(center, eye, up)

    image = renderer.render_to_image()
    return np.asarray(image)


def transform_mesh_vertices(vertices, rotation, translation, scale):
    if isinstance(vertices, np.ndarray):
        vertices = torch.tensor(vertices, dtype=torch.float32)

    vertices = vertices.unsqueeze(0)

    # Pre-transforms (Keep these as they are)
    vertices = vertices @ R_flip_z.to(vertices.device)
    vertices = vertices @ R_yup_to_zup.to(vertices.device)

    # Learned Transform
    R_mat = quaternion_to_matrix(rotation.to(vertices.device))
    tfm = Transform3d(dtype=vertices.dtype, device=vertices.device)
    tfm = (
        tfm.scale(scale)
        .rotate(R_mat)
        .translate(translation[0], translation[1], translation[2])
    )
    vertices_world = tfm.transform_points(vertices)

    # --- 2. Apply the new Z-Up Matrix here ---
    vertices_world = vertices_world @ R_pytorch3d_to_zup.to(vertices_world.device)

    return vertices_world[0]


def main():
    parser = argparse.ArgumentParser(
        description="Sample SAM3D object from image and mask, then export aligned OBJ."
    )
    parser.add_argument("--image", required=True, type=str, help="Path to input image")
    parser.add_argument(
        "--mask", required=True, type=str, help="Path to input mask file"
    )
    parser.add_argument(
        "--output_folder", required=True, type=str, help="Path to output folder"
    )
    parser.add_argument(
        "--seed", type=int, default=42, help="Random seed for inference"
    )

    args = parser.parse_args()

    args.image = Path(args.image)
    args.mask = Path(args.mask)
    args.output_folder = Path(args.output_folder)

    config_path = Path(
        ".checkpoints/hf/pipeline.yaml"
    )
    inference = Inference(str(config_path), compile=False)

    image = load_image(str(args.image))

    print(f"Loading mask: {args.mask}")
    mask_pil = Image.open(args.mask).convert("L")
    mask = np.array(mask_pil) > 0  # Create boolean mask

    output = inference(image, mask, seed=args.seed)

    mesh = output["glb"]
    vertices = mesh.vertices

    S = output["scale"][0].cpu().float()
    T = output["translation"][0].cpu().float()
    R = output["rotation"].squeeze().cpu().float()

    vertices_transformed = transform_mesh_vertices(vertices, R, T, S)

    mesh_transformed = copy.deepcopy(mesh)
    mesh_transformed.vertices = vertices_transformed.cpu().numpy().astype(np.float32)

    args.output_folder.mkdir(parents=True, exist_ok=True)

    mesh_output_path = args.output_folder / "mesh.obj"
    mesh_transformed.export(str(mesh_output_path))

    H, W = image.shape[0], image.shape[1]
    device = "cuda"
    image_tensor = (
        torch.from_numpy(np.array(image)).float().permute(2, 0, 1) / 255.0
    )
    image_tensor = image_tensor.to(device)
    moge_output = get_moge_pointcloud(image_tensor, device)
    K_matrix = denormalize_intrinsics(moge_output["intrinsics"].cpu().numpy(), H, W)

    intrinsics_output_path = args.output_folder / "intrinsics.npy"
    np.save(intrinsics_output_path, K_matrix)

    # Render
    img_array = render_mesh_with_intrinsics(str(mesh_output_path), K_matrix, W, H)
    render_image = Image.fromarray(img_array)
    render_output_path = args.output_folder / "render.png"
    render_image.save(render_output_path)

    img_array_pil = Image.fromarray(image)
    overlayed = Image.blend(img_array_pil.convert("RGBA"), render_image.convert("RGBA"), alpha=0.5)
    overlay_output_path = args.output_folder / "overlay.png"
    overlayed.save(overlay_output_path)





if __name__ == "__main__":
    main()

[snowflakewang]

@RafailFridman Hello, thank you for your great effort in writing the rendering code! I wonder why you select Mesh assets, instead of 3DGS ones, to render images?

Thank you!

[RafailFridman]

@snowflakewang In my project I want to use meshes later, and not 3DGS, that's it :)
I saw this issue #75 , but for me it is not critical

[snowflakewang]

@RafailFridman Thank you for the kind reply! I will try it!

[zh9586]

@RafailFridman谢谢你的回复！我会试试的！

Did you have any success with that approach? If so, would you be willing to share the code?