Diagnosing objectnav dataset from habitat-lab

diagnose_objectnav_dataset.py

import os
import glob
import json
import gzip
import habitat_sim


DATASET_ROOT = "data/datasets/objectnav/mp3d/v1/train"
SCENE_DATASET_ROOT = "data/scene_datasets/mp3d"


def make_habitat_configuration(mesh_path):
    # simulator configuration
    backend_cfg = habitat_sim.SimulatorConfiguration()
    backend_cfg.scene.id = mesh_path

    # agent configuration
    sensor_cfg = habitat_sim.SensorSpec()
    sensor_cfg.uuid = "semantic"
    sensor_cfg.resolution = [1080, 960]
    sensor_cfg.sensor_type = habitat_sim.SensorType.SEMANTIC
    agent_cfg = habitat_sim.agent.AgentConfiguration()
    agent_cfg.sensor_specifications = [sensor_cfg]

    return habitat_sim.Configuration(backend_cfg, [agent_cfg])


def robust_load_sim(scene_path):
    sim_cfg =  make_habitat_configuration(scene_path)
    hsim = habitat_sim.Simulator(sim_cfg)
    if not hsim.pathfinder.is_loaded:
        navmesh_settings = habitat_sim.NavMeshSettings()
        navmesh_settings.set_defaults()
        hsim.recompute_navmesh(hsim.pathfinder, navmesh_settings)
    return hsim


def verify_scene_episodes(scene_name):
    scene_path = f"{SCENE_DATASET_ROOT}/1pXnuDYAj8r/{scene_name}.glb"
    episodes_path = f"{DATASET_ROOT}/content/{scene_name}.json.gz"
    metadata_path = glob.glob(f"{DATASET_ROOT}/*.json.gz")[0]
    hsim = robust_load_sim(scene_path)
    with gzip.open(episodes_path, "rt") as fp:
        episodes_data = json.load(fp)
        goals_by_category = episodes_data["goals_by_category"]
    objects_in_scene = {}
    for level in hsim.semantic_scene.levels:
        for region in level.regions:
            for obj in region.objects:
                if obj.category.name() not in objects_in_scene:
                    objects_in_scene[obj.category.name()] = []
                obj_info = {
                    "center": obj.aabb.center,
                    "id": obj.id,
                }
                objects_in_scene[obj.category.name()].append(obj_info)
                
    # Match object ids in scene and in episode
    for ep_cat_name, ep_cat_info in goals_by_category.items():
        cat_name = ep_cat_name[len(scene_name + ".glb_"):]
        ep_obj_ids = set([ep["object_id"] for ep in ep_cat_info])
        scene_obj_ids = set([
            int(obj["id"].split("_")[-1]) for obj in objects_in_scene[cat_name]
        ])
        diff_ids = scene_obj_ids.difference(ep_obj_ids)
        if len(diff_ids) > 0:
            print(f"Object name: {cat_name} | Missing ids: {diff_ids}")


def main():
    scene_name = "1pXnuDYAj8r"
    verify_scene_episodes(scene_name)

    
if __name__ == "__main__":
    main()