refactor: 将子模块转为普通目录，移除外部 git 依赖

- 移除 JEPA/lejepa-identifiability 子模块 gitlink
- 移除 research/multiply/MultiPLY 子模块 gitlink
- 删除 .gitmodules（不再有外部 URL 依赖）
- 两个目录内容作为普通文件纳入主仓库追踪
- 删除各自内部 .git 目录，消除嵌套 git 仓库

research/multiply/MultiPLY/simulator/build_grid.py

@@ -0,0 +1,41 @@
+import os
+import quaternion # Remove this will cause invalid pointer error !!!!
+import habitat_sim
+import numpy as np
+from utils import config
+
+
+class GridBuilder(habitat_sim.Simulator):
+    def __init__(self, scene):
+        self.scene = scene
+        backend_cfg = habitat_sim.SimulatorConfiguration()
+        
+        backend_cfg.scene_id = os.path.join(config.HM3D_DIR, scene, f"{scene.split('-')[1]}.basis.glb")
+        # TODO: change this
+        backend_cfg.scene_dataset_config_file = os.path.join(config.HM3D_DIR, "hm3d_annotated_train_basis.scene_dataset_config.json")
+        backend_cfg.load_semantic_mesh = True
+        backend_cfg.enable_physics = False
+        cfg = habitat_sim.Configuration(backend_cfg, [habitat_sim.agent.AgentConfiguration()])
+        super().__init__(cfg)
+
+    def build_grids_if_not_exist(self):
+        _num = self.pathfinder.num_islands
+        _area = [self.pathfinder.island_area(x) for x in range(_num)]
+        _idx = _area.index(max(_area)) # Assert only one largest island
+        vertices = self.pathfinder.build_navmesh_vertices(_idx)
+        unique_vertices = np.unique(vertices, axis=0)
+
+        save_path = os.path.join(config.SAMPLE_DIR, self.scene)
+        os.makedirs(save_path, exist_ok=True)
+        np.save(os.path.join(save_path, "grid_points.npy"), unique_vertices)
+        # Nearest point in cube (cKDTree)
+
+
+if __name__ == "__main__":
+    print("Switch to latest version of habitat before running the code! Or Error")
+    folder_list = [folder for folder in  os.listdir(config.HM3D_DIR) if folder.startswith("00") and len(os.listdir(os.path.join(config.HM3D_DIR, folder))) == 4]
+
+    for i in os.listdir(config.HM3D_DIR):
+        if os.path.isdir(os.path.join(config.HM3D_DIR, i)):
+            sim = GridBuilder(i)
+            sim.build_grids_if_not_exist()

research/multiply/MultiPLY/simulator/feature_extractor.py

@@ -0,0 +1,198 @@
+import os
+import cv2
+import json
+import itertools
+import quaternion # Remove this will cause invalid pointer error !!!!
+import habitat_sim
+import numpy as np
+from tqdm import tqdm
+from utils import config
+from utils.dataset_interface import Objaverse, HM3D, ObjectFolder
+from multisensory_simulator import MultisensorySimulator
+from utils.config import sim_conf
+from utils.cloud_point_utils import Reconstruct3D
+from model.feature_encoder import LlaVa_Encoder
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data import DataLoader
+from PIL import Image
+import copy
+from collections import defaultdict
+import random
+
+class GridSampler(MultisensorySimulator):
+    def __init__(self, scene, new_objs=None, audio=False, encoder=None):
+        action_space = {
+            "look_left": habitat_sim.agent.ActionSpec("look_left", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_up": habitat_sim.agent.ActionSpec("look_up", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_down": habitat_sim.agent.ActionSpec("look_down", habitat_sim.agent.ActuationSpec(amount=90.0))}
+        cfg = sim_conf(scene, audio=audio)
+        cfg.agents[0].action_space.update(action_space)
+
+        super().__init__(cfg, new_objs)
+
+        self.scene = scene
+        _spec = cfg.agents[0].sensor_specifications[0]
+        self.reconstructor = Reconstruct3D(
+            _spec.resolution[0],
+            _spec.resolution[1],
+            float(_spec.hfov),
+            _spec.position
+        )
+        self.encoder = encoder
+
+    @staticmethod
+    def inside_p(pt, box):
+        if pt[0] < box[0][0] or pt[0] > box[1][0]: return False
+        if pt[1] < box[0][1] or pt[1] > box[1][1]: return False
+        if pt[2] < box[0][2] or pt[2] > box[1][2]: return False
+        if box[1][0] - pt[0] < 0 or pt[0] - box[0][0] < 0: return False
+        if box[1][2] - pt[2] < 0 or pt[2] - box[0][2] < 0: return False
+        return True
+
+    def scan_scene(self, bbox_file, return_features = True):
+        scene = bbox_file.split("_")[0]+".json"
+        room = bbox_file.replace(".json", "").split("_")[1]
+        room_bbox = json.load(open(os.path.join(config.ROOM_BBOX_DIR, scene)))[room]
+        room_bbox = [[room_bbox[0][0], room_bbox[0][2], room_bbox[0][1]], [room_bbox[1][0], room_bbox[1][2], room_bbox[1][1]]]
+        grid_points = np.load(os.path.join(config.SAMPLE_DIR, self.scene, "grid_points.npy"))
+        quats = [self.degree2quat(*x) for x in [(0, 0), (90, 0), (180, 0), (270, 0), (0, 90), (0, -90)]]
+
+        points = []
+        i = 0
+
+        all_instance_feature_dict = defaultdict(list)
+        all_instance_feature_dict_final = dict()
+
+        for x in tqdm([x for x in grid_points if self.pathfinder.is_navigable(x)]):
+            # if i > 10: continue
+            if not self.inside_p(x, room_bbox): continue
+            new_state = habitat_sim.AgentState(x, [0, 0, 0, 1])
+            self.agents[0].set_state(new_state) # set position
+
+            # Scan in sphere
+            obs = [self.parse_visual_observation(self.get_sensor_observations()),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left"))]
+
+            self.agents[0].set_state(new_state)  # reset
+            obs.append(self.parse_visual_observation(self.step("look_up")))
+            self.agents[0].set_state(new_state)  # reset
+            obs.append(self.parse_visual_observation(self.step("look_down")))
+
+            if return_features:
+                instance_feature_dict = self.get_per_instance_feature(i, bbox_file.replace(".json", ""), obs)
+                for instance, feature in instance_feature_dict.items():
+                    all_instance_feature_dict[instance].append(feature)
+
+            i += 6
+
+        return all_instance_feature_dict_final
+
+    def get_per_instance_feature(self, i, room, obs):
+        instance_feature_dict = dict()
+
+        for (j,frame) in enumerate(obs):
+            image = frame[..., :3].astype(np.uint8)  
+            image_features = self.encoder.encode(image)
+            pil_image = Image.fromarray(image)
+
+
+            all_semantics = frame[..., 4]
+            semantics = np.unique(all_semantics).astype(int)
+
+            for semantic in semantics:
+                # if semantic < 10000: continue
+                indices = np.where(all_semantics == semantic)
+                if indices[0].shape[0] < 10: continue
+                ymin, ymax, xmin, xmax = np.min(indices[0]), np.max(indices[0]), np.min(indices[1]), np.max(indices[1])
+                image_copy = copy.deepcopy(image)
+                image_copy[all_semantics != semantic] = 255
+                pil_image = Image.fromarray(image_copy)
+                
+                # 
+
+                cropped_image = pil_image.crop((xmin-1, ymin-1, xmax+1, ymax+1))
+                cropped_image = np.array(cropped_image)
+
+                pil_image.save("./tmp/%s/%d_%d.jpg"%(room, semantic, i+j))
+
+                cropped_features = self.encoder.encode(cropped_image)
+                instance_feature_dict[semantic] = cropped_features.mean(1).detach().cpu().numpy()
+
+        return instance_feature_dict
+
+    def parse_visual_observation(self, obs):
+        rgb = cv2.cvtColor(obs["rgba"], cv2.COLOR_RGBA2RGB)
+        frame = np.concatenate([rgb, obs["depth"][:, :, np.newaxis], obs["semantic"][:, :, np.newaxis]], axis=-1)
+
+        return frame
+
+    def get_semantic_labels(self):
+        id2cate = dict()
+        if self.new_objs is not None:
+            for i in self.new_objs:
+                id2cate[i["semantic_id"]] = i["cate"]
+        with open(os.path.join(config.HM3D_DIR, _scene, f"{_scene.split('-')[1]}.semantic.txt"), "r") as f:
+            a = f.readlines()
+        for i in a[1:]:
+            i = i.strip()
+            if len(i):
+                _id = int(i.split(",")[0])
+                _cate = i.split(",")[2].strip('"')
+                id2cate[_id] = _cate
+        return id2cate
+
+    @staticmethod
+    def degree2quat(z=0, x=0):
+        assert (z * x) == 0
+        if z:
+            half_radians = np.deg2rad(z) / 2.0
+            around_z_axis = [0, np.sin(half_radians), 0, np.cos(half_radians)]  # anticlockwise
+            return around_z_axis
+        elif x:
+            half_radians = np.deg2rad(x) / 2.0
+            around_x_axis = [np.sin(half_radians), 0, 0, np.cos(half_radians)]  # up is positive direction
+            return around_x_axis
+        else:
+            return [0, 0, 0, 1]
+
+
+if __name__ == "__main__":
+    objaverse = Objaverse(selected=False)
+    objectfolder = ObjectFolder()
+    bbox_dir = config.BBOX_WITH_ADDED_OBJECTS_DIR
+
+    # for bbox_file in os.listdir(bbox_dir):
+    for bbox_file in ["00009-vLpv2VX547B_7.json"]:
+        print ("Processing %s"%bbox_file)
+        bboxes = json.load(open(os.path.join(bbox_dir, bbox_file)))["incremented_bboxes"]
+        new_objs = []
+        objectfolder_cats = []
+        scene = bbox_file.split("_")[0]
+        objaverse_dict = dict()
+        for bbox in bboxes:
+
+            if "source" in bbox and bbox["source"] == "objaverse":
+                class_name = bbox["class_name"].replace("(soft)", "").replace("(hard)", "").replace("(deformable)", "").replace("(not deformable)", "").strip()
+                
+                try:    
+                    if class_name in objaverse_dict: id2 = objaverse_dict[class_name]
+                    else: id2 = random.choice(objaverse.lvis[class_name]); objaverse_dict[class_name] = id2
+                    path = objaverse.get_objects([id2])[id2]
+                    new_obj = {"path": path, "bbox": bbox["bbox"], "id": bbox["id"]}
+                    new_objs.append(new_obj)                 
+                except:
+                    continue
+
+        encoder = LlaVa_Encoder()
+        sampler = GridSampler(scene, new_objs, audio=False, encoder=encoder)
+        print ("successfully building sampler")
+
+        room = bbox_file.replace(".json", "").split("_")[1]
+        room_bbox = json.load(open(os.path.join(config.ROOM_BBOX_DIR, bbox_file.split("_")[0]+".json")))[room]
+        room_bbox = [[room_bbox[0][0], room_bbox[0][2], room_bbox[0][1]], [room_bbox[1][0], room_bbox[1][2], room_bbox[1][1]]]
+
+        feature_dict = sampler.scan_scene(bbox_file, return_features=True)
+        
+        print (feature_dict)

research/multiply/MultiPLY/simulator/grape_object.py

@@ -0,0 +1,101 @@
+import math
+import numpy as np
+import magnum as mn
+import habitat_sim
+from simulator.multisensory_simulator import MultisensorySimulator
+from utils.config import sim_conf
+from utils.dataset_interface import Objaverse
+from habitat_sim.utils import viz_utils as vut
+from habitat_sim.utils.common import quat_rotate_vector, quat_to_magnum
+
+
+class GrapeObject(MultisensorySimulator):
+    def __init__(self, scene, new_objs=None):
+        cfg = sim_conf(scene, audio=False, physics=True)
+        super().__init__(cfg, new_objs)
+        self.fetchable_objs = [x["obj_id"] for x in self.new_objs if "mass" in x]
+        self.reachable_range = 1.
+        self.reachable_degree = math.radians(180)
+        self.rigid_obj_mgr = self.get_rigid_object_manager()
+        self.fetched_obj = None
+        self.obj2agent = 0.3
+        super().step_physics(dt=2) # Let obj fail
+
+    def fetch_object(self, obj_id: int) -> bool:
+        if self.fetched_obj: return False
+        if not self.check_in_range(obj_id): return False
+        self.fetched_obj = self.rigid_obj_mgr.get_object_by_id(obj_id)
+        self.fetched_obj.motion_type = habitat_sim.physics.MotionType.KINEMATIC
+        self._update_fetched_obj()
+        return True
+
+    def _update_fetched_obj(self):
+        self.fetched_obj.translation = self.agent_center + quat_rotate_vector(self.agent_rot, [0, 0, -1]) * self.obj2agent
+        self.fetched_obj.rotation = quat_to_magnum(self.agent_rot)
+
+    def check_in_range(self, obj_id: int) -> bool:
+        # Check obj exist and fetchable
+        if obj_id not in self.fetchable_objs: return False
+        obj = self.rigid_obj_mgr.get_object_by_id(obj_id)
+        if obj is None: return False
+
+        # Check obj in dist range
+        obj_loc = obj.translation
+        dist = np.linalg.norm(obj_loc - self.agent_center)
+        print("Distance", dist, obj_loc)
+        if dist > self.reachable_range: return False
+
+        # Check obj in angle range
+        normal_vector = quat_rotate_vector(self.agent_rot, [0, 0, -1])
+        relative_vector = obj_loc - self.agent_center
+        dot = normal_vector[0] * relative_vector[0] + normal_vector[2] * relative_vector[2]
+        det = normal_vector[0] * relative_vector[2] - normal_vector[2] * relative_vector[0]
+        angle = math.atan2(det, dot)
+        print("Angle", angle, normal_vector, relative_vector)
+        if abs(angle) > (self.reachable_degree / 2): return False
+        return True
+
+    @property
+    def agent_center(self):
+        return self.agent_loc + [0, self.agents[0].agent_config.height - 0.3, 0]
+
+    def drop_object(self, drop_dist=0.) -> bool:
+        if not self.fetched_obj: return False
+        self.fetched_obj.motion_type = habitat_sim.physics.MotionType.DYNAMIC
+        self.fetched_obj.translation += quat_rotate_vector(self.agent_rot, [0, 0, -1]) * drop_dist
+        self.fetched_obj = None
+        return True
+
+    # Override step to update fetched object states
+    def step(self, action, dt=0.016666666666666666):
+        super().step(action, dt)
+        if self.fetched_obj:
+            self._update_fetched_obj()
+        return self.get_sensor_observations()
+
+    def step_physics(self, dt: float, scene_id: int = 0) -> None:
+        super().step_physics(dt, scene_id)
+        self.observations.append(self.get_sensor_observations())
+
+
+if __name__ == "__main__":
+    objaverse = Objaverse()
+    _objs = [
+        {"cate": "donut", "bbox": [[-6.0, -1.2, 1.0], [-6.0, -1.2, 1.0 + 0.075]],
+         "obj": "dcb0d1c9b8be49e0945535fdd81c7525", "mass": 0.05},
+    ]
+    for i in _objs:
+        i["path"] = objaverse.get_objects([i["obj"]])[i["obj"]]
+    sim = GrapeObject('00800-TEEsavR23oF', _objs)
+
+    sim.move_agent_to_target([-6.73648, 0.163378, -1.21183])
+    for i in range(3):
+        sim.step_physics(1. / 10)
+    success = sim.fetch_object(sim.fetchable_objs[0])
+    print(success)
+
+    sim.move_agent_to_target([-0.797001, 0.163378, -2.39349])
+    sim.drop_object(drop_dist=0.4)
+    for i in range(20):
+        sim.step_physics(1. / 10)
+    vut.make_video(sim.observations, "rgba", "color", "../color.mp4", fps=10, open_vid=False)

research/multiply/MultiPLY/simulator/grid_sampler.py

@@ -0,0 +1,190 @@
+import os
+import cv2
+import json
+import itertools
+import quaternion # Remove this will cause invalid pointer error !!!!
+import habitat_sim
+import numpy as np
+from tqdm import tqdm
+from utils import config
+from utils.dataset_interface import Objaverse, HM3D
+from simulator.multisensory_simulator import MultisensorySimulator
+from utils.config import sim_conf
+from utils.cloud_point_utils import Reconstruct3D, crop_points
+
+
+class GridSampler(MultisensorySimulator):
+    def __init__(self, scene, new_objs=None, audio=False, rooms=None):
+        self.scene = scene
+        self.grid_points = np.load(os.path.join(config.SAMPLE_DIR, self.scene, "grid_points.npy"))
+        self.room_boxes = []
+        if rooms is not None:
+            hm3d = HM3D()
+            for i in rooms:
+                bboxes = hm3d.load_room(f"{scene}_{i}.json")
+                _min, _max = hm3d.room_bbox(bboxes)
+                _min = [_min[0], _min[2], _min[1]]
+                _max = [_max[0], _max[2], _max[1]]
+                self.room_boxes.append([_min, _max])
+            self.grid_points = crop_points(self.grid_points, self.room_boxes)
+            print(self.room_boxes)
+
+        action_space = {
+            "look_left": habitat_sim.agent.ActionSpec("look_left", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_up": habitat_sim.agent.ActionSpec("look_up", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_down": habitat_sim.agent.ActionSpec("look_down", habitat_sim.agent.ActuationSpec(amount=90.0))}
+        cfg = sim_conf(scene, audio=audio)
+        cfg.agents[0].action_space.update(action_space)
+        super().__init__(cfg, new_objs)
+
+        _spec = cfg.agents[0].sensor_specifications[0]
+        self.reconstructor = Reconstruct3D(
+            _spec.resolution[0],
+            _spec.resolution[1],
+            float(_spec.hfov),
+            _spec.position
+        )
+
+    # Return:
+    # agent_location -> grid_points.npy
+    # camera_direction -> [left0, left90, left180, left270, up90, down90]
+    def scan_scene(self):
+        quats = [self.degree2quat(*x) for x in [(0, 0), (90, 0), (180, 0), (270, 0), (0, 90), (0, -90)]]
+
+        points = []
+        for i in tqdm([x for x in self.grid_points if self.pathfinder.is_navigable(x)]):
+            new_state = habitat_sim.AgentState(i, [0, 0, 0, 1])
+            self.agents[0].set_state(new_state) # set position
+
+            # Scan in sphere
+            obs = [self.parse_visual_observation(self.get_sensor_observations()),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left"))]
+
+            self.agents[0].set_state(new_state)  # reset
+            obs.append(self.parse_visual_observation(self.step("look_up")))
+            self.agents[0].set_state(new_state)  # reset
+            obs.append(self.parse_visual_observation(self.step("look_down")))
+
+            # convert to points
+            coordinates = []
+            valid_masks = []
+            for o, q in zip(obs, quats):
+                p, valid_mask = self.reconstructor.depth_map2points(o[:, :, 3], q, i)
+                coordinates.append(p)
+                valid_masks.append(valid_mask)
+            coordinates = np.concatenate(coordinates, axis=0)
+            valid_idx = np.where(np.concatenate(valid_masks, axis=0))[0]
+
+            # Concat all
+            obs = np.stack(obs, axis=0).reshape(-1, 5)
+            _points = np.concatenate([coordinates, obs[:, :3], obs[:, 4:]], axis=1).astype(np.float16) # xzy, rgb, semantic
+            _points = _points[valid_idx]
+            _points = _points[np.random.choice(len(_points), int(len(_points) / 10), replace=False), :]
+            points.append(_points)
+
+        points = np.concatenate(points, axis=0)
+        if len(self.room_boxes):
+            points = crop_points(points, self.room_boxes)
+        idx = self.reconstructor.downsample_index(points[:, :3])
+        return points[idx, :]
+
+    @staticmethod
+    # Channels -> [r, g, b, depth, semantic]
+    def parse_visual_observation(obs):
+        rgb = cv2.cvtColor(obs["rgba"], cv2.COLOR_RGBA2RGB)
+        frame = np.concatenate([rgb, obs["depth"][:, :, np.newaxis], obs["semantic"][:, :, np.newaxis]], axis=-1)
+        return frame
+
+    def get_semantic_labels(self):
+        id2cate = dict()
+        if self.new_objs is not None:
+            for i in self.new_objs:
+                id2cate[i["semantic_id"]] = i["cate"]
+        with open(os.path.join(config.HM3D_DIR, self.scene, f"{self.scene.split('-')[1]}.semantic.txt"), "r") as f:
+            a = f.readlines()
+        for i in a[1:]:
+            i = i.strip()
+            if len(i):
+                _id = int(i.split(",")[0])
+                _cate = i.split(",")[2].strip('"')
+                id2cate[_id] = _cate
+        return id2cate
+
+    @staticmethod
+    def degree2quat(z=0, x=0):
+        assert (z * x) == 0
+        if z:
+            half_radians = np.deg2rad(z) / 2.0
+            around_z_axis = [0, np.sin(half_radians), 0, np.cos(half_radians)]  # anticlockwise
+            return around_z_axis
+        elif x:
+            half_radians = np.deg2rad(x) / 2.0
+            around_x_axis = [np.sin(half_radians), 0, 0, np.cos(half_radians)]  # up is positive direction
+            return around_x_axis
+        else:
+            return [0, 0, 0, 1]
+
+# def sample_rirs(table, fps):
+#     hm3d = HM3D()
+#     top_center = lambda x: [(x[0][0] + x[1][0])/2, x[1][2], (x[0][1] + x[1][1])/2]
+#     for _scene, v in table.items():
+#         new_objs = list(itertools.chain.from_iterable(v.values()))
+#         existing_objs = hm3d.load_scene(_scene)
+#         obj2loc = {x["obj"]: top_center(x["bbox"]) for x in new_objs}
+#         obj2loc.update({x["id"]: top_center(x["bbox"]) for x in existing_objs})
+#
+#         cfg = sim_conf(_scene, visual=False, audio=True)
+#         sim = MultisensorySimulator(cfg, fps=fps, new_objs=None)
+#         rirs = dict()
+#         for k, v in obj2loc.items():
+#             sim.set_audio_source(v)
+#             obs = grid_sampling(sim, _scene)
+#
+#             _rirs = []
+#             for i in obs:
+#                 _r = []
+#                 for j in i:
+#                     _r.append(j["audio_sensor"])
+#                 _rirs.append(_r)
+#             rirs[k] = _rirs
+#         json.dump(rirs, open(os.path.join(config.SAMPLE_DIR, _scene, "rirs.json"), "w"))
+
+
+if __name__ == "__main__":
+    # TODO: audio sampler after task template & grid_point navigable - coord dict to rirs
+    # TODO: - change loop in calculate_audio() and reverb from previous time step
+    objaverse = Objaverse()
+    scene = json.load(open(os.path.join(config.DATA_DIR, "scene.json"), "r"))
+    objaverse.get_objects([x["obj"] for x in scene]) # Download objects
+
+    table = dict()
+    for i in scene:
+        _scene = i["room"].split("_")[0]
+        if _scene not in table: table[_scene] = dict()
+        _trail = i["trail"]
+        if _trail not in table[_scene]: table[_scene][_trail] = []
+        i["path"] = objaverse.get_objects([i["obj"]])[i["obj"]]
+        table[_scene][_trail].append(i)
+
+    for _scene, v in table.items():
+        for _trail, _objs in v.items():
+            print(_scene, _trail, _objs)
+
+            sampler = GridSampler(_scene, _objs, audio=False)
+            # Semantic Labels
+            _path = os.path.join(config.SAMPLE_DIR, _scene, f"{_trail}.json")
+            json.dump(sampler.get_semantic_labels(), open(_path, "w"))
+
+            # Sampling
+            points = sampler.scan_scene()
+            np.save(os.path.join(config.SAMPLE_DIR, _scene, f"{_trail}.npy"), points)
+            # Visualize (reverse y)
+            # with open(os.path.join(config.SAMPLE_DIR, _scene, f"{_trail}.txt"), "w") as file:
+            #     file.write(f"{len(points)}\n")
+            #     for p in points:
+            #         file.write(f"{p[0]} {-p[2]} {p[1]} {p[3]} {p[4]} {p[5]}\n")
+
+    # Not support material yet: https://github.com/facebookresearch/sound-spaces/issues/111
+    # sim.set_material_file("audio_sensor", "data/HM3D/mp3d_material_config.json")

research/multiply/MultiPLY/simulator/multisensory_simulator.py

@@ -0,0 +1,222 @@
+import quaternion # Remove this will cause invalid pointer error !!!!
+import habitat_sim
+import habitat
+from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
+from habitat.sims.habitat_simulator.actions import HabitatSimActions
+
+import os
+import librosa
+import numpy as np
+import magnum as mn
+import matplotlib.pyplot as plt
+from scipy.signal import fftconvolve
+from moviepy.editor import VideoFileClip
+from habitat_sim.utils import viz_utils as vut
+from moviepy.audio.AudioClip import AudioArrayClip
+from utils import config
+from typing import List
+
+os.environ['MAGNUM_LOG'] = "quiet"
+os.environ['HABITAT_SIM_LOG'] = "quiet"
+
+
+class MultisensorySimulator(habitat_sim.Simulator):
+    def __init__(self, conf: habitat_sim.Configuration, new_objs: List[dict] = None):
+        super().__init__(conf)
+        # Fake habitat.core.simulator.Simulator
+        self._sim = self
+        # self.habitat_config = habitat.get_config()
+        # self.habitat_config["SCENE"] = conf.sim_cfg.scene_id
+
+        # Assign semantic ids & place new objs
+        self.new_objs = new_objs
+        if self.new_objs is not None:
+            count = 0
+            for i in self.new_objs:
+                _id = 10000 + count
+                i["semantic_id"] = _id
+                count += 1
+            self._place_objs()
+
+        # Others
+        assert len(self.agents) == 1
+        
+        # TODO: change this
+        self.observations = [self.get_sensor_observations()] # Init observation after audio setup
+
+    def _place_objs(self):
+        obj_attr_mgr = self.get_object_template_manager()
+        rigid_obj_mgr = self.get_rigid_object_manager()
+
+        for i in self.new_objs: # v in (x, y, z) format but hm3d in (x, z, y) format
+            k = i["path"]
+            v = i["bbox"]
+            # Calc scale
+            object_template = obj_attr_mgr.create_new_template(k)
+            obj_temp_id = obj_attr_mgr.register_template(object_template)
+            obj = rigid_obj_mgr.add_object_by_template_id(obj_temp_id)
+            _bbox = obj.root_scene_node.compute_cumulative_bb()
+            _scale = (v[1][2] - v[0][2]) / (_bbox.top - _bbox.bottom)
+            rigid_obj_mgr.remove_object_by_id(obj.object_id)
+            obj_attr_mgr.remove_template_by_id(obj_temp_id)
+
+            # Add new mesh
+            object_template.scale = np.ones(3) * _scale
+            object_template.semantic_id = i["semantic_id"]
+            obj_temp_id = obj_attr_mgr.register_template(object_template)
+            obj = rigid_obj_mgr.add_object_by_template_id(obj_temp_id)
+            i["obj_id"] = obj.object_id
+
+            # Move object
+            _loc = [(v[0][0] + v[1][0]) / 2, v[0][2], (v[0][1] + v[1][1]) / 2]
+
+            _bbox = obj.root_scene_node.compute_cumulative_bb()
+            obj.translation = -_bbox.center() + _loc
+            if "rot" in i:
+                obj.rotation = mn.Quaternion.rotation(mn.Deg(i["rot"]), [0.0, 1.0, 0.0])
+            if "mass" in i:
+                obj.motion_type = habitat_sim.physics.MotionType.DYNAMIC
+                obj.mass = i["mass"]
+            else:
+                obj.motion_type = habitat_sim.physics.MotionType.STATIC
+
+            print(i["cate"], obj.translation, _scale)
+
+        # TODO: enable
+        # self.update_navmesh()
+
+    def update_navmesh(self):
+        # # recompute the NavMesh with STATIC objects
+        navmesh_settings = habitat_sim.NavMeshSettings()
+        navmesh_settings.set_defaults()
+        navmesh_settings.include_static_objects = True
+        navmesh_success = self.recompute_navmesh(self.pathfinder, navmesh_settings)
+        if not navmesh_success:
+            raise Exception("Recompute Navmesh Fail.")
+
+    def set_audio_source(self, loc):
+        audio_sensor = self.get_agent(0)._sensors["audio_sensor"]
+        audio_sensor.setAudioSourceTransform(loc)
+
+    def show_top_down_map(self, meters_per_pixel=0.1, height=0., path_points=None):
+        print(f"The NavMesh bounds in {height} are: " + str(self.pathfinder.get_bounds()))
+        top_down_map = self.pathfinder.get_topdown_view(meters_per_pixel, height)
+        top_down_map = 1. - top_down_map * 0.5 # Recolor
+
+        plt.figure(figsize=(12, 8))
+        plt.axis("off")
+        plt.imshow(top_down_map, cmap='gray', vmin=0., vmax=1.)
+
+        top_down_loc = self._convert_points_to_topdown([self.agent_loc], meters_per_pixel)[0]
+        plt.plot(*top_down_loc, marker="o", markersize=10, alpha=0.8)
+
+        if path_points:
+            top_down_loc = self._convert_points_to_topdown(path_points, meters_per_pixel)
+            plt.plot(*np.array(top_down_loc).transpose(), marker="o", markersize=5, alpha=0.8)
+        plt.show()
+
+    def _convert_points_to_topdown(self, points, meters_per_pixel):
+        bounds = self.pathfinder.get_bounds()
+
+        # convert 3D x,z to topdown x,y
+        points_topdown = []
+        for point in points:
+            px = (point[0] - bounds[0][0]) / meters_per_pixel
+            py = (point[2] - bounds[0][2]) / meters_per_pixel
+            points_topdown.append(np.array([px, py]))
+        return points_topdown
+
+    def _path_planning(self, target_loc):
+        path = habitat_sim.ShortestPath()
+        path.requested_start = self.agent_loc
+        path.requested_end = target_loc
+        found_path = self.pathfinder.find_path(path)
+        if found_path:
+            return path.points
+        else:
+            return [self.agent_loc]
+
+    def move_agent_to_target(self, target_loc, goal_radius=1., final_goal_radius=0.):
+        # First point is current position. Last point is target location.
+        path_points = self._path_planning(target_loc)
+        print(f"Move path {path_points}")
+
+        shortest_path_follower = ShortestPathFollower(sim=self, goal_radius=goal_radius, return_one_hot=False)
+        for idx, i in enumerate(path_points):
+            if (idx + 1) == len(path_points):
+                shortest_path_follower = ShortestPathFollower(sim=self, goal_radius=final_goal_radius, return_one_hot=False)
+            while True:
+                next_action = shortest_path_follower.get_next_action(i)
+                if next_action == HabitatSimActions.stop:
+                    break
+                elif next_action == HabitatSimActions.move_forward:
+                    action = "move_forward"
+                elif next_action == HabitatSimActions.turn_left:
+                    action = "turn_left"
+                elif next_action == HabitatSimActions.turn_right:
+                    action = "turn_right"
+                else:
+                    raise Exception(f"Action {next_action} not defined.")
+
+                assert action in self.agent_actions.keys()
+                obs = self.step(action)
+                self.observations.append(obs)
+
+    def calculate_audio(self):
+        rirs = [np.array(x["audio_sensor"]).T for x in self.observations]
+        audio_data, _ = librosa.load(self.audio_objs[0].audio_path, sr=config.RIR_SAMPLING_RATE)
+
+        index = 0
+        audio = []
+        # TODO: change this
+        scaled_sample_rate = int(config.RIR_SAMPLING_RATE * self.step_delta_t)
+        for i in rirs:
+            if index * scaled_sample_rate - i.shape[0] < 0:
+                source_sound = audio_data[: (index + 1) * scaled_sample_rate]
+                binaural_convolved = np.array([fftconvolve(source_sound, i[:, channel]) for channel in range(i.shape[-1])])
+                audio_goal = binaural_convolved[:, index * scaled_sample_rate: (index + 1) * scaled_sample_rate]
+            else:
+                # include reverb from previous time step
+                source_sound = audio_data[index * scaled_sample_rate - i.shape[0] + 1: (index + 1) * scaled_sample_rate]
+                binaural_convolved = np.array([fftconvolve(source_sound, i[:, channel], mode='valid') for channel in range(i.shape[-1])])
+                audio_goal = binaural_convolved
+            audio.append(audio_goal)
+            index = (index + 1) % (audio_data.shape[0] // scaled_sample_rate)
+        return np.concatenate(audio, axis=-1).transpose()
+
+    def set_material_file(self, sensor_key, file_path):
+        self.agents[0]._sensors[sensor_key].setAudioMaterialsJSON(file_path)
+
+    @property
+    def agent_loc(self):
+        return self.agents[0].state.position
+
+    @property
+    def agent_rot(self):
+        return self.agents[0].state.rotation
+
+    @property
+    def agent_actions(self):
+        return dict(self.agents[0].agent_config.action_space)
+
+
+# TODO: move all audio related code into fake sim
+def demo(sim, fps, target=None):
+    # Navigation and first-person video
+    if target is None:
+        target = sim.pathfinder.get_random_navigable_point()
+
+    sim.move_agent_to_target(target)
+    vut.make_video(sim.observations, "rgba", "color", "../color.mp4", fps=fps, open_vid=False)
+    vut.make_video(sim.observations, "depth", "depth", "../depth.mp4", fps=fps, open_vid=False)
+    vut.make_video(sim.observations, "semantic", "semantic", "../semantic.mp4", fps=fps, open_vid=False)
+
+    # Calculate audio and merge with video
+    # audio = sim.calculate_audio()
+    # _audio = AudioArrayClip(audio, fps=RIR_SAMPLING_RATE)
+    # _audio.write_audiofile("./demo.wav")
+
+    # _video = VideoFileClip("./tmp.mp4")
+    # _video = _video.set_audio(_audio)
+    # _video.write_videofile("./demo.mp4")
+    # os.remove("tmp.mp4")

research/multiply/MultiPLY/simulator/semantic_extractor.py

@@ -0,0 +1,260 @@
+import os
+import cv2
+import json
+import itertools
+import quaternion # Remove this will cause invalid pointer error !!!!
+import habitat_sim
+import numpy as np
+from tqdm import tqdm
+from utils import config
+from utils.dataset_interface import Objaverse, HM3D, ObjectFolder
+from multisensory_simulator import MultisensorySimulator
+from utils.config import sim_conf
+from utils.cloud_point_utils import Reconstruct3D
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data import DataLoader
+from PIL import Image
+import copy
+from collections import defaultdict
+import random
+from utils.cloud_point_utils import Reconstruct3D, crop_points
+
+class GridSampler(MultisensorySimulator):
+    def __init__(self, scene, new_objs=None, audio=False):
+        action_space = {
+            "look_left": habitat_sim.agent.ActionSpec("look_left", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_up": habitat_sim.agent.ActionSpec("look_up", habitat_sim.agent.ActuationSpec(amount=90.0)),
+            "look_down": habitat_sim.agent.ActionSpec("look_down", habitat_sim.agent.ActuationSpec(amount=90.0))}
+        cfg = sim_conf(scene, audio=audio)
+        cfg.agents[0].action_space.update(action_space)
+
+        super().__init__(cfg, new_objs)
+
+        self.scene = scene
+        _spec = cfg.agents[0].sensor_specifications[0]
+        self.reconstructor = Reconstruct3D(
+            _spec.resolution[0],
+            _spec.resolution[1],
+            float(_spec.hfov),
+            _spec.position
+        )
+        # self.encoder = encoder
+
+    @staticmethod
+    def inside_p(pt, box):
+        if pt[0] < box[0][0] or pt[0] > box[1][0]: return False
+        if pt[1] < box[0][1] or pt[1] > box[1][1]: return False
+        if pt[2] < box[0][2] or pt[2] > box[1][2]: return False
+        if box[1][0] - pt[0] < 0 or pt[0] - box[0][0] < 0: return False
+        if box[1][2] - pt[2] < 0 or pt[2] - box[0][2] < 0: return False
+        return True
+
+    def scan_scene(self, bbox_file, all_ids, scene, return_features = True):
+        room_bbox = bbox_file
+
+        grid_points = np.load(os.path.join(config.SAMPLE_DIR, self.scene, "grid_points.npy"))
+
+
+        grid_points = crop_points(grid_points, [room_bbox])
+        quats = [self.degree2quat(*x) for x in [(0, 0), (90, 0), (180, 0), (270, 0), (0, 90), (0, -90)]]
+
+        points = []
+        i = 0
+
+        all_instance_feature_dict = defaultdict(list)
+        all_instance_feature_dict_final = dict()
+
+        for x in [x for x in grid_points if self.pathfinder.is_navigable(x)]:    
+            new_state = habitat_sim.AgentState(x, [0, 0, 0, 1])
+            self.agents[0].set_state(new_state) # set position
+
+            # Scan in sphere
+            obs = [self.parse_visual_observation(self.get_sensor_observations()),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left")),
+                   self.parse_visual_observation(self.step("look_left"))]
+
+
+            self.get_per_instance(i, scene.replace(".json", ""), obs, all_ids)
+
+            i += 6
+
+            # convert to points
+            coordinates = []
+            valid_masks = []
+            for o, q in zip(obs, quats):
+                p, valid_mask = self.reconstructor.depth_map2points(o[:, :, 3], q, x)
+                coordinates.append(p)
+                valid_masks.append(valid_mask)
+            coordinates = np.concatenate(coordinates, axis=0)
+            valid_idx = np.where(np.concatenate(valid_masks, axis=0))[0]
+
+            # Concat all
+            obs = np.stack(obs, axis=0).reshape(-1, 5)
+            _points = np.concatenate([coordinates, obs[:, :3], obs[:, 4:]], axis=1).astype(np.float16) # xzy, rgb, semantic
+            _points = _points[valid_idx]
+            points.append(_points)
+        
+
+        if not len(points): return np.zeros((1,1))
+
+        print ("%d vertices inside the room"%len(points))
+        points = np.concatenate(points, axis=0)
+        print ("%d points inside the room"%points.shape[0])
+        points = crop_points(points, [room_bbox])
+        print ("%d points after crop"%points.shape[0])
+        if points.shape[0] == 0: return np.zeros((1,1))
+
+        idx = self.reconstructor.downsample_index(points[:, :3])
+        points = points[idx, :]
+        print ("%d points after sampling"%points.shape[0])
+        
+        return points
+
+    def get_per_instance(self, i, room, obs, all_ids):
+        instance_feature_dict = dict()
+
+        for (j,frame) in enumerate(obs):
+            image = frame[..., :3].astype(np.uint8)  
+            depth = frame[..., 3].astype(np.uint8)  
+
+            pil_image = Image.fromarray(image)
+
+            try:
+                os.mkdir("./data/original_2d_gt_seg/%s"%room)
+            except:
+                pass
+
+            np.save("./data/original_2d_gt_seg/%s/depth_%d.npy"%(room, i+j), depth)
+            pil_image.save("./data/original_2d_gt_seg/%s/image_%d.jpg"%(room, i+j))
+
+            all_semantics = frame[..., 4]
+            semantics = np.unique(all_semantics).astype(int)
+
+            for semantic in semantics:
+                if not (semantic in all_ids or str(semantic) in all_ids or semantic >= 10000): continue
+                
+                indices = np.where(all_semantics == semantic)
+                if np.min(indices[0]) == 0 or np.min(indices[1]) == 0 or np.max(indices[1]) == 719 or np.max(indices[0]) == 719: continue
+                if indices[0].shape[0] < 100: continue 
+                
+                ymin, ymax, xmin, xmax = np.min(indices[0]), np.max(indices[0]), np.min(indices[1]), np.max(indices[1])
+                image_copy = copy.deepcopy(image)
+                image_copy[all_semantics != semantic] = 255
+                pil_image = Image.fromarray(image_copy)
+                
+                pil_image.save("./data/original_2d_gt_seg/%s/%d_%d.jpg"%(room, semantic, i+j))
+
+                cropped_image = pil_image.crop((xmin-1, ymin-1, xmax+1, ymax+1))
+
+                cropped_image.save("./data/original_2d_gt_seg/%s/%d_%d_cropped.jpg"%(room, semantic, i+j))
+
+    # Add features
+    def parse_visual_observation(self, obs):
+        rgb = cv2.cvtColor(obs["rgba"], cv2.COLOR_RGBA2RGB)
+        frame = np.concatenate([rgb, obs["depth"][:, :, np.newaxis], obs["semantic"][:, :, np.newaxis]], axis=-1)
+
+        return frame
+
+    def get_semantic_labels(self):
+        id2cate = dict()
+        if self.new_objs is not None:
+            for i in self.new_objs:
+                id2cate[i["semantic_id"]] = i["cate"]
+        with open(os.path.join(config.HM3D_DIR, _scene, f"{_scene.split('-')[1]}.semantic.txt"), "r") as f:
+            a = f.readlines()
+        for i in a[1:]:
+            i = i.strip()
+            if len(i):
+                _id = int(i.split(",")[0])
+                _cate = i.split(",")[2].strip('"')
+                id2cate[_id] = _cate
+        return id2cate
+
+    @staticmethod
+    def degree2quat(z=0, x=0):
+        assert (z * x) == 0
+        if z:
+            half_radians = np.deg2rad(z) / 2.0
+            around_z_axis = [0, np.sin(half_radians), 0, np.cos(half_radians)]  # anticlockwise
+            return around_z_axis
+        elif x:
+            half_radians = np.deg2rad(x) / 2.0
+            around_x_axis = [np.sin(half_radians), 0, 0, np.cos(half_radians)]  # up is positive direction
+            return around_x_axis
+        else:
+            return [0, 0, 0, 1]
+
+
+if __name__ == "__main__":
+    objaverse = Objaverse(selected=False)
+    objectfolder = ObjectFolder()
+
+    bbox_dir = config.HM3D_BBOX_DIR
+    hm3d = HM3D()
+
+    for bbox_file in tqdm(os.listdir(bbox_dir)):
+
+        print ("Processing %s"%bbox_file)
+        bboxes = json.load(open(os.path.join(bbox_dir, bbox_file)))
+ 
+        new_objs = []
+        objectfolder_cats = []
+        scene = bbox_file.split("_")[0]
+        objaverse_dict = dict()
+
+        final_bboxes = []
+        original_bboxes = []
+        all_ids = []
+
+        bbox_file_copy = bbox_file
+        scene2 = bbox_file.split("_")[0]+".json"
+        room2 = bbox_file.replace(".json", "").split("_")[1]
+        try:
+            room_bbox = json.load(open(os.path.join(config.ROOM_BBOX_DIR, scene2)))[room2]
+        except:
+            continue
+
+        k = 1
+
+        for bbox in bboxes:
+            if "source" in bbox and bbox["source"] == "objectfolder":
+                try:
+                    cat, id2, material, path = objectfolder.get_objects(bbox["class_name"].replace("_hot", "").replace("_cold", ""))
+                    new_obj = {"path": path, "bbox": bbox["bbox"], "id": bbox["id"]}
+                    new_objs.append(new_obj)
+                    final_bboxes.append(bbox)
+                except:
+                    continue
+
+            elif "source" in bbox and bbox["source"] == "objaverse":
+                class_name = bbox["class_name"].replace("_hard", "").replace("_soft", "").replace("_hot", "").replace("_cold", "").strip()
+                
+                try:    
+                    id2 = random.choice(objaverse.lvis[class_name])
+                    path = objaverse.get_objects([id2])[id2]
+                    new_obj = {"path": path, "bbox": bbox["bbox"], "id": bbox["id"]}
+                    new_objs.append(new_obj) 
+                    final_bboxes.append(bbox)                
+                except:
+                    continue
+
+            else:
+                if not bbox['class_name'] in ['floor', 'wall', 'ceiling']:
+                    all_ids.append(bbox['id'])
+                    final_bboxes.append(bbox)
+
+                original_bboxes.append(bbox)
+            
+            final_bboxes.append(bbox)
+        
+
+        sampler = GridSampler(scene, new_objs, audio=False)
+        print ("successfully building sampler")
+
+        points = sampler.scan_scene(room_bbox, all_ids, bbox_file_copy, return_features=True)
+        sampler.close()
+
+        if points.shape[0] == 1: continue
+
+        np.save(os.path.join(config.SAMPLE_DIR, scene, f"{bbox_file_copy}.npy"), points)

research/multiply/MultiPLY/simulator/utils/config.py

@@ -0,0 +1,83 @@
+import os.path
+from pathlib import Path
+
+
+# Dir
+ROOT_DIR = Path(__file__).parent.parent
+DATA_DIR = os.path.join(ROOT_DIR, "data")
+HM3D_DIR = os.path.join(DATA_DIR, "hm3d")
+HM3D_BBOX_DIR = os.path.join(DATA_DIR, "hm3d_obj_bbox")
+OBJAVERSE_DIR = os.path.join(DATA_DIR, "objaverse")
+AUDIOSET_DIR = os.path.join(DATA_DIR, "audio_set")
+TASK_TEMPLATE_DIR = os.path.join(DATA_DIR, "task_template")
+SAMPLE_DIR = os.path.join(DATA_DIR, "sampled_data")
+OBJECTFOLDER_DIR = os.path.join(DATA_DIR, "object_folder")
+OBJECTFOLDER_OBJECTS_DIR = os.path.join(DATA_DIR, "ObjectFolder")
+BBOX_WITH_ADDED_OBJECTS_DIR = os.path.join(DATA_DIR, "bbox_with_added_objects")
+BBOX_WITH_ADDED_OBJECTFOLDER_DIR = os.path.join(DATA_DIR, "bbox_with_added_objectfolder")
+BBOX_WITH_TEMPERATURE_DIR = os.path.join(DATA_DIR, "bbox_with_temperature")
+ROOM_BBOX_DIR = os.path.join(DATA_DIR, "room_bboxes")
+THIRD_PARTY_DIR = os.path.join(ROOT_DIR, "third_party")
+
+# GPT
+OPENAI_KEY = ""
+# OPENAI_PROXY = {"http": "127.0.0.1:7890", "https": "127.0.0.1:7890"}
+OPENAI_PROXY = {}
+
+# Simulator
+RIR_SAMPLING_RATE = 16000
+def sim_conf(scene: str, visual=True, audio=True):
+    import quaternion # Remove this will cause invalid pointer error !!!!
+    import habitat_sim
+    backend_cfg = habitat_sim.SimulatorConfiguration()
+    backend_cfg.scene_id = os.path.join(HM3D_DIR, scene, f"{scene.split('-')[1]}.basis.glb")
+    # TODO: change this
+    backend_cfg.scene_dataset_config_file = os.path.join(HM3D_DIR, "hm3d_annotated_train_basis.scene_dataset_config.json")
+    backend_cfg.load_semantic_mesh = True
+    backend_cfg.enable_physics = False
+
+    sensors = []
+    if visual:
+        camera_resolution = [720, 720] # h = w for scene scan
+        camera_position = [0.0, 1.4, 0.0]
+        _spec = habitat_sim.CameraSensorSpec()
+        _spec.uuid = "rgba"
+        _spec.sensor_type = habitat_sim.SensorType.COLOR
+        _spec.resolution = camera_resolution
+        _spec.position = camera_position
+        _spec.orientation = [0.0, 0.0, 0.0]
+        _spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE
+        sensors.append(_spec)
+
+        _spec = habitat_sim.CameraSensorSpec()
+        _spec.uuid = "depth"
+        _spec.sensor_type = habitat_sim.SensorType.DEPTH # COLOR = 1, DEPTH = 2, SEMANTIC = 4
+        _spec.resolution = camera_resolution
+        _spec.position = camera_position
+        _spec.orientation = [0.0, 0.0, 0.0]
+        _spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE
+        sensors.append(_spec)
+
+        _spec = habitat_sim.CameraSensorSpec()
+        _spec.uuid = "semantic"
+        _spec.sensor_type = habitat_sim.SensorType.SEMANTIC
+        _spec.resolution = camera_resolution
+        _spec.position = camera_position
+        _spec.orientation = [0.0, 0.0, 0.0]
+        _spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE
+        sensors.append(_spec)
+
+    if audio:
+        _spec = habitat_sim.AudioSensorSpec()
+        _spec.uuid = "audio_sensor" # Must use this name or backend simulator will raise error :(
+        _spec.enableMaterials = False
+        _spec.channelLayout.type = habitat_sim.sensor.RLRAudioPropagationChannelLayoutType.Binaural
+        _spec.channelLayout.channelCount = 2
+        _spec.acousticsConfig.sampleRate = RIR_SAMPLING_RATE
+        _spec.acousticsConfig.indirect = True
+        sensors.append(_spec)
+
+    agent_cfg = habitat_sim.agent.AgentConfiguration()
+    agent_cfg.sensor_specifications = sensors
+    cfg = habitat_sim.Configuration(backend_cfg, [agent_cfg])
+    return cfg

research/multiply/MultiPLY/simulator/utils/dataset_interface.py

@@ -0,0 +1,498 @@
+import copy
+import os
+import librosa
+import random
+import numpy as np
+import pandas as pd
+import json
+from typing import List
+import itertools
+from collections import defaultdict
+from utils import config
+import re
+from tqdm import tqdm
+import objaverse
+
+class HM3D(object):
+    def __init__(self):
+        self.dir_path = config.HM3D_BBOX_DIR
+        self.augmented_dir_path = config.BBOX_WITH_ADDED_OBJECTS_DIR
+        self.objectfolder_dir_path = config.BBOX_WITH_ADDED_OBJECTFOLDER_DIR
+        self.scene2cate = dict()
+        self.files = sorted(x for x in os.listdir(self.dir_path) if len(self.load_room(x)) > 0)
+        self.augmented_files = sorted(x for x in os.listdir(self.augmented_dir_path) if len(self.load_room(x)) > 0)
+        self.objectfolder_files = sorted(x for x in os.listdir(self.objectfolder_dir_path) if len(self.load_room(x)) > 0)
+
+        for i in self.files:
+            scene = i.split("_")[0]
+            if scene not in self.scene2cate:
+                self.scene2cate[scene] = []
+
+            bboxes = json.load(open(os.path.join(self.dir_path, i), "r"))
+            for b in bboxes:
+                self.scene2cate[scene].append(b["class_name"])
+        self.categories = list(sorted(set(itertools.chain.from_iterable(self.scene2cate.values()))))
+        self.scenes = list(sorted(set(self.scene2cate.keys())))
+
+    def load_room(self, room_json: str):
+        # Coordinates in (x,y,z) format
+        bboxes = json.load(open(os.path.join(self.dir_path, room_json), "r"))
+        return bboxes
+
+    def load_room_with_added_objects(self, room_json: str):
+        # Coordinates in (x,y,z) format
+        bboxes = json.load(open(os.path.join(self.augmented_dir_path, room_json), "r"))
+        return bboxes
+
+    def load_room_with_added_objectfolder(self, room_json: str):
+        # Coordinates in (x,y,z) format
+        bboxes = json.load(open(os.path.join(self.objectfolder_dir_path, room_json), "r"))
+        return bboxes
+
+    def load_scene(self, scene: str):
+        bboxes = []
+        for i in self.files:
+            if scene in i:
+                _b = self.load_room(i)
+                bboxes.extend(_b)
+        return bboxes
+
+    def scene_description(self, room_json: str, max_box=50, format="topbottom"):
+        bboxes = self.load_room(room_json)
+        if len(bboxes) > max_box:
+            bboxes = random.choices(bboxes, k=max_box)
+        np.set_printoptions(suppress=True)
+        _min = np.round(np.min(np.array([x["bbox"][0] for x in bboxes]), axis=0), 3)
+        _max = np.round(np.max(np.array([x["bbox"][1] for x in bboxes]), axis=0), 3)
+        
+        if format == "topbottom":
+            room_desc = f"<room>: [{list(_min)}, {list(_max)}]\n"
+            obj_desc = "\n".join([f'<{x["class_name"]}>({x["id"]}): {np.round(x["bbox"], 3)}'.replace("\n", ",") for x in bboxes])
+            return room_desc + obj_desc
+        elif format == "center":
+            obj_desc = "\n".join([f'<{x["class_name"]}>({x["id"]}): {np.round((np.array(x["bbox"][0]) + np.array(x["bbox"][1])) / 2, 3)}'.replace("\n", ",") for x in bboxes])
+            return obj_desc
+        elif format == "object_name":
+            obj_desc = "\n".join([f'<{x["class_name"]}>({x["id"]})'.replace("\n", ",") for x in bboxes])
+            return obj_desc
+
+    def multi_modal_scene_description(self, room_json: str, max_box=50):
+        try:
+            bboxes = self.load_room_with_added_objects(room_json)['incremented_bboxes']
+        except:
+            bboxes = self.load_room_with_added_objectfolder(room_json)['incremented_bboxes']
+        bboxes = bboxes[:max_box]
+        np.set_printoptions(suppress=True)
+
+        final_bboxes = []
+        for bbox in bboxes:
+            if "source" in bbox and bbox["source"] == "objaverse":
+                bbox["class_name"] = bbox["class_name"] + "(audio, tactile)"
+            elif "source" in bbox and bbox["source"] == "objectfolder":
+                bbox["class_name"] = bbox["class_name"] + "(tapsound)"
+            new_bbox_format = f'<{bbox["class_name"]}>({bbox["id"]}): {np.round(bbox["bbox"], 3)}'.replace("\n", ",")
+
+            final_bboxes.append(new_bbox_format)
+
+        obj_desc = "\n".join(final_bboxes)
+        return obj_desc
+
+
+# TODO: include evaluation set
+class AudioSet(object):
+    def __init__(self, training_set=True):
+        self.dir_path = config.AUDIOSET_DIR
+
+        # Load
+        ontology = json.load(open(os.path.join(self.dir_path, "ontology.json"), "r"))
+        if training_set:
+            meta_file = "unbalanced_train_segments.csv"
+        else:
+            meta_file = "eval_segments.csv"
+        meta = pd.read_csv(os.path.join(self.dir_path, meta_file), sep=", ", engine='python', skiprows=2)
+
+        # Set selected categories
+        # Ontology is a graph, not a tree. query_handles is visible tree roots.
+        query_handles = ["Music", "Sounds of things"]
+        valid_cate = ["Musical instrument", "Domestic sounds, home sounds", "Liquid", "Glass", "Printer",
+                      "Air conditioning", "Mechanical fan", "Clock", "Fire alarm", "Smoke detector, smoke alarm",
+                      "Doorbell", "Alarm clock", "Ringtone", "Telephone bell ringing", "Domestic sounds, home sounds",
+                      "Loudspeaker", "Radio", "Television", "MP3", "Domestic animals, pets"]
+        block_cate = ["Human sounds", "Vehicle"]
+
+        # Put audios on the node.
+        name2node = {x["name"]: x["id"] for x in ontology}
+        node2child = {x["id"]: x["child_ids"] for x in ontology}
+        valid_nodes = self.iterative_query([name2node[x] for x in valid_cate], query_dict=node2child)
+        block_nodes = self.iterative_query([name2node[x] for x in block_cate], query_dict=node2child)
+
+        self._node2audio = defaultdict(list)
+        for id, labels in zip(meta["# YTID"], meta["positive_labels"]):
+            labels = set(labels.strip('"').split(","))
+            if len(labels & block_nodes): continue
+            for i in (labels & valid_nodes):
+                self._node2audio[i].append(id)
+
+        # Pruning nodes without audios
+        self.nodes = list()
+        for i in [x["id"] for x in ontology]:
+            nodes = self.iterative_query([i], node2child)
+            if any(len(self._node2audio[x]) for x in nodes):
+                self.nodes.append(i)
+
+        query_nodes = self.iterative_query([name2node[x] for x in query_handles], query_dict=node2child)
+        self.nodes = list(set(self.nodes) & query_nodes)
+
+        filtered_ontology = [x for x in ontology if x["id"] in self.nodes]
+        self.node2name = {x["id"]: x["name"] for x in filtered_ontology}
+        self.node2description = {x["id"]: x["description"] for x in filtered_ontology}
+        self.node2child = {x["id"]: (set(x["child_ids"]) & set(self.nodes)) for x in filtered_ontology}
+        self.node2father = defaultdict(list)
+        for k, v in self.node2child.items():
+            for i in v:
+                self.node2father[i].append(k)
+
+        # Others
+        self.audio_ids = self.get_ids(self.nodes)
+        self.meta = meta[meta["# YTID"].isin(self.audio_ids)]
+        self.downloader = os.path.join(config.THIRD_PARTY_DIR, "youtube-dl")
+        print(f"AudioSet {meta_file}: {len(self.meta)} / {len(meta)}, cate {len(self.nodes)} / {len(ontology)}")
+
+        # _str = filtered.to_csv(index=False, sep="\t") # Stupid Lib
+        # _str = _str.replace("\t", ", ")
+        # with open(os.path.join(self.dir_path, f"filtered_{meta_file}"), "w") as f:
+        #     f.write(_str)
+
+        # Display
+        # root_nodes = set(self.nodes).difference(set(itertools.chain.from_iterable(self.node2child.values())))
+        # self.print_tree(root_nodes)
+
+    def get_ids(self, nodes: List[str]):
+        nodes = self.iterative_query(nodes, self.node2child)
+        return list(set(itertools.chain.from_iterable(self._node2audio[x] for x in nodes)))
+
+    def get_audio(self, audio_id):
+        assert audio_id in self.audio_ids # YTID is unique in training set
+        info = self.meta[self.meta["# YTID"] == audio_id]
+        assert len(info) == 1
+        _path = os.path.join(config.AUDIOSET_DIR, f"{audio_id}.wav")
+        if not os.path.exists(_path):
+            os.system(f"sh {os.path.join(config.THIRD_PARTY_DIR, 'fetch_audio.sh')} "
+                      f"{audio_id} {info['start_seconds'].values[0]} {info['end_seconds'].values[0]} "
+                      f"{_path} {self.downloader}")
+
+        audio_data = None
+        success = False
+        if os.path.exists(_path):
+            audio_data, _ = librosa.load(_path, sr=config.RIR_SAMPLING_RATE)
+            success = True
+        return audio_data, success
+
+    @staticmethod
+    def iterative_query(nodes: List[str], query_dict: dict[str, List[str]], include_root=True) -> set:
+        q = copy.deepcopy(nodes)
+        res = []
+        while len(q):
+            node = q.pop()
+            res.append(node)
+            q.extend(query_dict[node])
+
+        if not include_root:
+            for i in nodes:
+                res.remove(i)
+        return set(res)
+
+    def print_tree(self, nodes, max_depth=100):
+        q = []
+        for i in nodes:
+            q.append((i, 0))
+        while len(q):
+            node, depth = q.pop()
+            for i in self.node2child[node]:
+                q.append((i, depth+1))
+            if depth < max_depth:
+                num = len(set(itertools.chain.from_iterable(
+                    self._node2audio[x] for x in self.iterative_query([node], self.node2child))))
+                print(f'{"--" * depth} {self.node2name[node]}: {num}, {self.node2description[node]}')
+
+    @property
+    def meta_info(self):
+        info = {}
+        for i in self.nodes:
+            path = self.iterative_query([i], self.node2father)
+            tags = [self.node2name[x] for x in path]
+            description = self.node2description[i]
+            info[self.node2name[i]] = f"tags={tags}, description='{description}'"
+        return info
+
+
+class Objaverse(object):
+    def __init__(self, selected=True):
+        self.dir_path = config.OBJAVERSE_DIR
+        if not selected:
+            with open(os.path.join(config.OBJAVERSE_DIR, "audio_objaverse.txt"), "r") as f:
+                valid_cate = f.readlines()
+        else:
+            with open(os.path.join(config.OBJAVERSE_DIR, "selected_objaverse_lvis.txt"), "r") as f:
+                valid_cate = f.readlines()
+        valid_cate = [x.strip("\n") for x in valid_cate]
+        self.lvis = {k.strip(): v for k, v in objaverse.load_lvis_annotations().items() if k in valid_cate}
+        self.categories = sorted(valid_cate)
+
+        # self.meta_info = []
+        # for k, v in self.anns.items():
+        #     info = {'id': k}
+        #
+        #     if len(v["name"]):
+        #         info["name"] = v["name"]
+        #     if k in self.lvis: # Precise labels
+        #         info["label"] = self.lvis[k]
+        #     if len(v["categories"]):
+        #         info["categories"] = [x['name'] for x in v['categories']]
+        #     if len(v["tags"]):
+        #         info["tags"] = [x['name'] for x in v['tags']]
+        #     if len(v["description"]):
+        #         info["description"] = v['description']
+        #     self.meta_info.append(str(json.dumps(info)))
+
+    @staticmethod
+    def get_objects(uids):
+        return objaverse.load_objects(uids=uids, download_processes=1)
+
+
+class Objaverse_Material(object):
+    def __init__(self):
+        self.dir_path = config.DATA_DIR
+        self.all_objaverse_materials = json.load(open(os.path.join(self.dir_path, "objaverse_random_obj_material_dict.json")))
+
+        self.all_objects = []
+        self.all_cats = []
+
+        idx = 0
+
+        for cat, materials in self.all_objaverse_materials.items():
+            for material in materials:
+                material['obj_id'] = idx
+                self.all_objects.append(material)
+                idx += 1
+
+            self.all_cats.append(cat)
+    
+    def get_random_objs(self, num: int) -> list:
+        chosen_cats = np.random.choice(self.all_cats, num)
+
+        chosen_objects = []
+        find_ambiguous = False
+
+        for cat in chosen_cats:
+            chosen_objects.extend([str(obj) for obj in self.all_objaverse_materials[cat]])
+            if len(self.all_objaverse_materials[cat]) >= 2:
+                find_ambiguous = True
+
+        while not find_ambiguous:
+            cat = np.random.choice(self.all_cats, 1)[0]
+            if len(self.all_objaverse_materials[cat]) >= 2:
+                find_ambiguous = True
+                chosen_objects.extend([str(obj) for obj in self.all_objaverse_materials[cat]])
+
+        return chosen_objects
+
+
+class Objaverse_Material2(object):
+    def __init__(self):
+        self.dir_path = config.DATA_DIR
+        self.all_objects = json.load(open(os.path.join(self.dir_path, "objaverse_random_obj_material_list.json")))
+
+    def get_random_objs(self, num: int) -> list:
+        index = random.randint(0, len(self.all_objects) - num)
+        chosen_objects = self.all_objects[index:index+num]
+        return chosen_objects
+
+
+class ObjectFolder(object):
+    def __init__(self):
+        self.dir_path = config.OBJECTFOLDER_DIR
+        self.obj2cate = dict()
+        meta = pd.read_csv(os.path.join(self.dir_path, "objects.csv"), header=None)
+        abo = pd.read_csv(os.path.join(self.dir_path, "abo_classes_3d.txt"), sep=",", header=None)
+        cate_map = dict(zip(meta[0].astype(int), meta[1]))
+        abo_map = dict(zip(abo[0], abo[1]))
+        self.id2material = dict(zip(meta[0].astype(int), meta[3]))
+
+        self.id2cate = dict()
+        for k, v in cate_map.items():
+            if v in abo_map:
+                v = abo_map[v]
+            self.id2cate[k] = v
+
+        self.categories = list(set(self.id2cate.values()))
+        cate2material = {x: [] for x in self.categories}
+        for k, v in self.id2cate.items():
+            cate2material[v].append(self.id2material[k])
+
+        select_cate = []
+        self.cate2materialset = dict()
+        for k, v in cate2material.items():
+            if len(set(v)) > 1:
+                self.cate2materialset[k] = list(set(v))
+                # print(k, set(v))
+                select_cate.append(k)
+        self.cate2ids = {x: [] for x in select_cate}
+        for k, v in self.id2cate.items():
+            if v in select_cate:
+                self.cate2ids[v].append(k)
+
+        # json.dump(cate2ids, open(os.path.join(self.dir_path, "cat2ids.json"), "w"))
+        # json.dump(cate2materialset, open(os.path.join(self.dir_path, "cat2mateiral.json"), "w"))
+        # json.dump(self.id2material, open(os.path.join(self.dir_path, "id2mateiral.json"), "w"))
+
+        # for i in os.listdir(self.dir_path):
+        #     if not os.path.isdir(os.path.join(self.dir_path, i)): continue
+        #     for j in os.listdir(os.path.join(self.dir_path, i)):
+        #         file_path = os.path.join(self.dir_path, i, j, "model.obj")
+        #         cate = cate_map[j]
+        #         if cate in abo_map: cate = abo_map[cate]
+        #         self.obj2cate[file_path] = cate
+        # self.categories = set(self.obj2cate.values())
+
+
+    # Before call this function, please download all ObjectFolder objects in the ObjectFolder directory
+
+    def get_objects(self, category):
+        material = re.findall("\((Iron|Wood|Plastic|Steel|Ceramic|Polycarbonate|Glass|iron|wood|plastic|steel|ceramic|polycarbonate|glass)\)", category)[0]
+
+        cat = category.replace(material, "").replace("(", "").replace(")", "").strip()
+        material = material.replace("(", "").replace(")", "")
+        
+        ids = self.cate2ids[cat]
+        final_ids = []
+
+        for id2 in ids:
+            if self.id2material[id2].lower() == material.lower():
+                final_ids.append(id2)
+
+        id2 = random.choice(final_ids)
+        path = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, str(id2), "model_new.obj")
+
+        return cat, id2, material, path
+
+    @staticmethod
+    def modify_obj(fn, new_fn):
+        fin = open(fn, 'r')
+        fout = open(new_fn, 'w')
+        lines = [line.rstrip() for line in fin]
+        fin.close()
+
+        vertices = []; normals = []; faces = []; vns = []
+        header = ""
+        for line in lines:
+            if line.startswith('v '):
+                vertice = np.float32(line.split()[1:4])
+                line = "v %f %f %f"%(vertice[0], vertice[2], vertice[1])
+
+            fout.write(line+"\n")
+        fout.close()
+
+    @staticmethod
+    def normalize_pts(pts):
+        out = np.array(pts, dtype=np.float32)
+        center = np.mean(out, axis=0)
+        out -= center
+        scale = np.sqrt(np.max(np.sum(out**2, axis=1)))
+        out /= scale
+        return out
+
+    @staticmethod
+    def load_obj(fn):
+        fin = open(fn, 'r')
+        lines = [line.rstrip() for line in fin]
+        fin.close()
+
+        vertices = []; normals = []; faces = [];
+        for line in lines:
+            if line.startswith('v '):
+                vertices.append(np.float32(line.split()[1:4]))
+            elif line.startswith('f '):
+                faces.append(np.int32([item.split('/')[0] for item in line.split()[1:4]]))
+
+        return vertices, faces
+
+    def rotate_and_normalize(self):
+        for obj in tqdm(os.listdir(config.OBJECTFOLDER_OBJECTS_DIR)):
+            model_file = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "model.obj")
+            new_model_file = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "model_new.obj")
+            print ("processing %s"%model_file)
+            self.modify_obj(model_file, new_model_file)
+
+    def generate_vertices_and_forces(self):
+        for obj in tqdm(os.listdir(config.OBJECTFOLDER_OBJECTS_DIR)):
+            model_file = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "model.obj")
+            save_vertice_file = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "vertices.npy")
+            save_force_file = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "forces.npy")
+            v, f = self.load_obj(model_file)
+            v = random.sample(v, 20)
+            forces = np.ones((20, 3))
+            v = np.vstack(v)
+            np.save(save_vertice_file, v); np.save(save_force_file, forces)
+
+    def embed_features(self):
+        from msclap import CLAP
+        import torch
+        from subprocess import call
+
+        clap_model = CLAP(version = '2023', use_cuda=False)
+
+        for obj in tqdm(os.listdir(config.OBJECTFOLDER_OBJECTS_DIR)):
+            try:
+                audio_dir = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "results")
+                feature_save_dir = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "features")
+                cmd = "rm -rf %s*"%feature_save_dir
+                call(cmd, shell=True)
+                os.mkdir(feature_save_dir)
+                if not os.path.exists(audio_dir):
+                    continue
+                
+                audio_files = os.listdir(audio_dir)
+                audio_files = [os.path.join(audio_dir, file) for file in audio_files]
+                audio_embeddings = clap_model.get_audio_embeddings(audio_files)
+                
+                for i in range(audio_embeddings.shape[0]):
+                    torch.save(audio_embeddings[i], feature_save_dir+"/"+str(i)+".pt")
+            except:
+                print ("failed processing clap features for %s" %obj)
+
+    def prepare_adapter_data(self):
+        from subprocess import call
+        question_dict = []
+
+        for obj in tqdm(os.listdir(config.OBJECTFOLDER_OBJECTS_DIR)):
+            feature_dir = os.path.join(config.OBJECTFOLDER_OBJECTS_DIR, obj, "features")
+            if not os.path.exists(feature_dir): continue
+            for (j,feature) in enumerate(os.listdir(feature_dir)):
+                try:
+                    os.mkdir("final_dataset/impact_sound_%s_%d"%(obj, j))
+                    os.mkdir("final_dataset/impact_sound_%s_%d/impact_sound"%(obj, j))
+                except:
+                    pass
+                cmd = "cp %s/%s final_dataset/impact_sound_%s_%d/impact_sound/0.pt"%(feature_dir, feature, obj, j)
+                call (cmd, shell=True)
+                question = "What's the material of the object? <Tap>"
+                answer = self.id2material[int(obj)]
+                question_dict.append({"impact_sound": "impact_sound_%s_%d"%(obj, j), "question": question, "answer": answer})
+
+        with open("questions/impact_sound_adapter.json", "w") as f:
+            json.dump(question_dict, f)
+
+if __name__ == "__main__":
+    # TODO: spilt train and test set (use src_file label for audio files)
+    # # hm3d = HM3D()
+    # # objaverse = Objaverse()
+    
+    # audio_set = AudioSet(training_set=True)
+    # node = random.choice(audio_set.nodes)
+    # cate_name = audio_set.node2name[node]
+    # audio_id = random.choice(audio_set.get_ids([node]))
+    # audio, success = audio_set.get_audio(audio_id)
+    # print(cate_name, audio_id, audio.shape, success)
+    objectfolder = ObjectFolder()
+    objectfolder.prepare_adapter_data()

research/multiply/MultiPLY/simulator/utils/reconstruct3d.py

@@ -0,0 +1,76 @@
+import math
+import random
+import numpy as np
+from scipy.spatial.transform import Rotation
+from tqdm import tqdm
+import open3d as o3d
+
+# TODO: rename to point_utils.py
+class Reconstruct3D(object):
+    def __init__(self, h, w, hfov, camera2agent):
+        self.h = h
+        self.w = w
+        self.focal_length = (w / 2) / math.tan(np.deg2rad(hfov / 2))
+        self.camera2agent = camera2agent
+        self.voxel_size = 0.5
+        self.num_points_per_voxel = 1000
+
+    def depth_map2points(self, depth_map, quat, loc):
+        rot = Rotation.from_quat(quat)
+
+        # Depth to agent coordinate
+        _max = 100
+        _min = 0
+        valid_mask = (depth_map > _min) & (depth_map < _max)
+        depth_map = np.clip(depth_map, _min, _max)
+
+        _x, _z = np.meshgrid(np.arange(self.w), np.arange(self.h - 1, -1, -1))
+        x = (_x - (self.w - 1) / 2.) * depth_map / self.focal_length
+        y = depth_map
+        z = (_z - (self.h - 1) / 2.) * depth_map / self.focal_length
+        _points = np.stack([x, z, y], axis=-1).reshape(-1, 3)
+        # Rotate points
+        _points = rot.inv().apply(_points)
+        # Agent to world coordinate
+        _points[:, 0] += loc[0]
+        _points[:, 1] += loc[1]
+        _points[:, 2] -= loc[2]  # reverse axis
+        return _points + self.camera2agent, valid_mask.reshape(-1)
+
+    def downsample_index(self, points):
+        # Drop far points
+        dist = np.linalg.norm(points, axis=1)
+        valid_idx = np.where(dist < 100)[0]
+        valid_points = points[valid_idx, :]
+
+        # Build voxels
+        min_coord = np.array([np.min(valid_points[:, 0]), np.min(valid_points[:, 1]), np.min(valid_points[:, 2])])
+        max_coord = np.array([np.max(valid_points[:, 0]), np.max(valid_points[:, 1]), np.max(valid_points[:, 2])])
+        num_voxels_x = int((max_coord[0] - min_coord[0]) / self.voxel_size) + 1
+        num_voxels_y = int((max_coord[1] - min_coord[1]) / self.voxel_size) + 1
+        num_voxels_z = int((max_coord[2] - min_coord[2]) / self.voxel_size) + 1
+        print(len(points), num_voxels_x, num_voxels_y, num_voxels_z)
+        voxel_grid = np.zeros((num_voxels_x, num_voxels_y, num_voxels_z), dtype=object)
+        for i in range(num_voxels_x):
+            for j in range(num_voxels_y):
+                for k in range(num_voxels_z):
+                    voxel_grid[i, j, k] = []
+
+        # Assign points to voxels
+        voxel_indices = ((valid_points - min_coord) / self.voxel_size).astype(int)
+        for i, idx in tqdm(zip(valid_idx, voxel_indices)):
+            voxel_grid[idx[0], idx[1], idx[2]].append(i)
+
+        # Random sampling in voxels
+        res = []
+        for i in tqdm(voxel_grid.flatten()):
+            if len(i) > self.num_points_per_voxel:
+                res.extend(np.random.choice(i, self.num_points_per_voxel, replace=False))
+            else:
+                res.extend(i)
+        return res
+
+    def crop_points(self, points, bbox):
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(points)
+        obj_points = open3d.geometry.crop_point_cloud(pcd, bbox[0], bbox[1])