tsdf_fusion_open3d.py

#!/usr/bin/env python

import numpy as np
import open3d
import trimesh

import objslampp


def get_scenes():
    with objslampp.utils.timer():
        dataset = objslampp.datasets.YCBVideoDataset('train')

    color_type = open3d.TSDFVolumeColorType.RGB8
    # color_type = open3d.TSDFVolumeColorType.Gray32
    # color_type = getattr(open3d.TSDFVolumeColorType, 'None')

    pipeline = open3d.ScalableTSDFVolume(
        voxel_length=0.0075,
        sdf_trunc=0.04,
        color_type=color_type,
    )
    image_ids = dataset.get_ids()

    scene_global = trimesh.Scene()
    for index in range(len(dataset)):
        image_id = image_ids[index]
        if image_id.split('/')[0] != '0000':
            break

        print(index)

        example = dataset.get_example(index)

        K = example['meta']['intrinsic_matrix']
        intrinsic = open3d.camera.PinholeCameraIntrinsic(
            height=example['depth'].shape[0],
            width=example['depth'].shape[1],
            fx=K[0, 0],
            fy=K[1, 1],
            cx=K[0, 2],
            cy=K[1, 2],
        )
        T_world2cam = np.r_[
            example['meta']['rotation_translation_matrix'],
            [[0, 0, 0, 1]],
        ]
        T_cam2world = np.linalg.inv(T_world2cam)
        color = open3d.Image(example['color'])
        depth = open3d.Image(example['depth'].astype(np.float32))
        rgbd = open3d.create_rgbd_image_from_color_and_depth(
            color=color,
            depth=depth,
            depth_trunc=1.0,
            depth_scale=1.0,
            convert_rgb_to_intensity=color_type != open3d.TSDFVolumeColorType.RGB8,  # NOQA
        )

        pipeline.integrate(
            image=rgbd,
            intrinsic=intrinsic,
            extrinsic=T_world2cam,
        )

        camera = trimesh.scene.Camera(
            resolution=(example['depth'].shape[1], example['depth'].shape[0]),
            focal=(K[0, 0], K[1, 1]),
            transform=T_cam2world,
        )
        geom = trimesh.creation.camera_marker(camera, marker_height=0.02)
        if index % 15 == 0:
            scene_global.add_geometry(geom)

        scene = scene_global.copy()
        scene.add_geometry(geom)

        mesh = pipeline.extract_triangle_mesh()
        mesh = objslampp.utils.open3d_to_trimesh(mesh)
        scene.add_geometry(mesh)

        scene.camera.transform = objslampp.extra.trimesh.to_opengl_transform(
            T_cam2world
        )
        yield {'image': example['color'], 'scene': scene}


objslampp.extra.trimesh.display_scenes(get_scenes(), tile=(1, 2))

tsdf_fusion_pykfusion.py

#!/usr/bin/env python

import imgviz
import numpy as np
import pykfusion

import objslampp


dataset = objslampp.datasets.YCBVideoDataset('train')

example = dataset.get_example(0)

rgb = example['color']
depth = (np.nan_to_num(example['depth']) * 1000).astype(np.uint16)
K = example['meta']['intrinsic_matrix']
T_world2cam = np.r_[
    example['meta']['rotation_translation_matrix'],
    [[0, 0, 0, 1]],
]
T_cam2world = np.linalg.inv(T_world2cam)

tsdf = pykfusion.pyKFusion(
    width=depth.shape[1],
    height=depth.shape[0],
    camera_intrinsic=K,
    voxel_size=0.0075,
    tsdf_dim=(512,) * 3,
    tsdf_pose=np.eye(4),
    extra_tsdf_dims=6,
    instance_id=0,
    image_buffers=True,
)
camera_pose, quat_jtj, error = tsdf.track(
    T_cam2world,
    np.eye(4),
    pykfusion.pyOpenCvUchar3(rgb),
    pykfusion.pyOpenCvUint16(depth),
    depth_weight=1.0,
    rgb_weight=0.0,
)
tsdf.integrate(
    camera_pose,
    rgb_image=pykfusion.pyOpenCvUchar3(rgb),
    depth_image=pykfusion.pyOpenCvUint16(depth),
    instance_id=0,
)
rgb_rendered = pykfusion.pyOpenCvUchar3(rgb)
depth_rendered = pykfusion.pyOpenCvUint16(depth)
pykfusion.render_tsdfs_cuda(
    tsdfs=[tsdf],
    pose=camera_pose,
    ids=[0],
    depths=depth_rendered,
    rgb=rgb_rendered,
)
rgb_rendered = np.asarray(rgb_rendered)[:, :, ::-1]
depth_rendered = np.asarray(depth_rendered)

tsdf_volume, tsdf_weight = tsdf.get_volume()

from skimage import measure
import trimesh

# Marching cubes
verts, faces, norms, vals = measure.marching_cubes_lewiner(tsdf_volume,level=0)
verts_ind = np.round(verts).astype(int)
# voxel grid coordinates to world coordinates
verts = trimesh.transform_points(verts * tsdf.voxel_size(), tsdf.tsdf_pose)

geom = trimesh.Trimesh(vertices=verts, faces=faces, vertex_normals=norms)
geom.show()


depth = depth.astype(np.float32) / 1000
depth[depth == 0] = np.nan
depth_rendered = depth_rendered.astype(np.float32) / 1000
depth_rendered[depth_rendered == 0] = np.nan

viz = imgviz.tile([
    rgb,
    rgb_rendered,
    imgviz.depth2rgb(depth, min_value=0.3, max_value=3),
    imgviz.depth2rgb(depth_rendered, min_value=0.3, max_value=3),
], border=(255,) * 3)
imgviz.io.pyglet_imshow(viz)
imgviz.io.pyglet_run()