Accurate global point clouds cannot be generated through depth sensors

[akaling15926]

Habitat-Lab and Habitat-Sim versions

Habitat-Lab: v0.3.3

Habitat-Sim: v0.3.3

The point cloud generated by my code has a lot of intersections and is chaotic, but I see that the scene point clouds in the dataset are very precise. I'm not sure if there's a problem with my code.

import os
import random
import numpy as np
import quaternion
import habitat
from habitat.config import read_write
from habitat.config.default import get_agent_config
import open3d as o3d

output_path = "./output"
os.makedirs(output_path, exist_ok=True)


config = habitat.get_config("third_party/habitat-lab/habitat-lab/habitat/config/benchmark/nav/objectnav/objectnav_mp3d.yaml")
with read_write(config):
    config.habitat.dataset.split = "val"
    agent_config = get_agent_config(config.habitat.simulator)
    config.habitat.simulator.turn_angle = 30
    agent_config.sim_sensors.depth_sensor.normalize_depth = False


hfov = float(agent_config.sim_sensors.depth_sensor.hfov) * np.pi / 180.

K = np.array([
    [1 / np.tan(hfov / 2.), 0., 0., 0.],
    [0., 1 / np.tan(hfov / 2.), 0., 0.],
    [0., 0., 1, 0],
    [0., 0., 0, 1]
])

W = agent_config.sim_sensors.depth_sensor.width
H = agent_config.sim_sensors.depth_sensor.height
def depth_to_pointcloud(depth, rotation, translation):
  
 
    xs, ys = np.meshgrid(np.linspace(-1, 1, W), np.linspace(1, -1, H))
    

    depth_flat = depth.reshape(1, -1)
    xs_flat = xs.reshape(1, -1)
    ys_flat = ys.reshape(1, -1)
    
 
    xys = np.vstack((xs_flat * depth_flat, 
                    ys_flat * depth_flat, 
                    -depth_flat, 
                    np.ones(depth_flat.shape)))
   
    xy_camera = np.matmul(np.linalg.inv(K), xys)
    

    rotation_matrix = quaternion.as_rotation_matrix(rotation)
    T_world_camera = np.eye(4)
    T_world_camera[0:3, 0:3] = rotation_matrix
    T_world_camera[0:3, 3] = translation
    

    points_world = np.matmul(T_world_camera, xy_camera)
    points_3d = points_world[0:3, :].T  # (N, 3)
    

    valid_mask = (depth_flat[0] > 0.1) & (depth_flat[0] < 10.0)
    
    return points_3d, valid_mask
def save_pointcloud(points, colors, filename='pointcloud.ply'):

    with open(filename, 'w') as f:
        f.write("ply\n")
        f.write("format ascii 1.0\n")
        f.write(f"element vertex {points.shape[0]}\n")
        f.write("property float x\n")
        f.write("property float y\n")
        f.write("property float z\n")
        f.write("property uchar red\n")
        f.write("property uchar green\n")
        f.write("property uchar blue\n")
        f.write("end_header\n")
        

        for i in range(points.shape[0]):
            f.write(f"{points[i, 0]} {points[i, 1]} {points[i, 2]} ")
            f.write(f"{int(colors[i, 0] * 255)} {int(colors[i, 1] * 255)} {int(colors[i, 2] * 255)}\n")
    
    print(f" {filename}")


env = habitat.Env(config=config)

env.episodes = random.sample(env.episodes, k=6)

for ep_idx in range(len(env.episodes)):
    print(f"\n=== Episode {ep_idx + 1} ===")
    obs = env.reset()
    print(f"Goal category: {env.current_episode.goals[0].object_category}")

    all_points = []
    all_colors = []

    max_steps = 13
    for step in range(max_steps):
        if env.episode_over:
            break

        depth = obs["depth"][...,0]
        rgb = obs["rgb"]
        camera_state = env._sim.get_agent_state()
        cam_rotation = camera_state.sensor_states['depth'].rotation
        cam_translation = camera_state.sensor_states['depth'].position

        points, valid_mask = depth_to_pointcloud(depth, cam_rotation, cam_translation)
        colors = rgb.reshape(-1, 3) / 255.0 
        
        valid_points = points[valid_mask]
        valid_colors = colors[valid_mask]
        all_points.append(valid_points)
        all_colors.append(valid_colors)

        obs = env.step(3)
    all_points = np.vstack(all_points)
    all_colors = np.vstack(all_colors)


    name = env.current_episode.scene_id.split('/')[-1].split('.')[0]

    output_filename = os.path.join(output_path, f"{name}_ep{ep_idx+1}_pointcloud.ply")
    save_pointcloud(None, all_points, all_colors, filename=output_filename)

env.close()