diff --git a/examples/viewer.py b/examples/viewer.py index 78df8231ba..3550ad4f5d 100644 --- a/examples/viewer.py +++ b/examples/viewer.py @@ -9,13 +9,17 @@ import sys import time from enum import Enum -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple, Union flags = sys.getdlopenflags() sys.setdlopenflags(flags | ctypes.RTLD_GLOBAL) import magnum as mn import numpy as np +from habitat.sims.habitat_simulator.sim_utilities import ( + get_all_objects, + get_obj_from_id, +) from magnum import shaders, text from magnum.platform.glfw import Application @@ -26,6 +30,209 @@ from habitat_sim.utils.settings import default_sim_settings, make_cfg +def find_interaction_surface_points( + sim: habitat_sim.Simulator, + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], + num_vertical_slices: int = 10, + num_radial_slices: int = 10, + cull_points=True, + max_point_set_size: int = 20, +) -> Tuple[List[mn.Vector3], List[mn.Vector3]]: + """ + Use raycasting to find a set of points on the lateral surfaces of an object. + + :return: the list of interaction points and the list of cast rays + """ + + assert num_vertical_slices >= 1, "Must at least slice in half." + assert num_radial_slices >= 4, "Must at least form a 2D simplex." + assert max_point_set_size >= 3, "Must at least form a 2D simplex." + + surface_points: List[mn.Vector3] = [] + + # compute the ray set: + aabb = obj.aabb + ray_set: List[habitat_sim.geo.Ray] = [] + + # compute the circle size to contain the aabb: xz diagonal length+10%. + # Used for max ray distance and cylinder sampling. + size_x = aabb.size_x() + size_y = aabb.size_y() + size_z = aabb.size_z() + circle_rad = math.sqrt((size_x / 2.0) ** 2 + (size_z / 2.0) ** 2) * 1.1 + if False: + # cylinder (furniture): + # for each vertical slice, select rays in a circle + for i in range(1, num_vertical_slices + 1): + y_val = aabb.bottom + (i / num_vertical_slices) * size_y + center_point = mn.Vector3(0, y_val, 0) + for r in range(num_radial_slices): + cx = circle_rad * math.cos(2 * math.pi * r / num_radial_slices) + cz = circle_rad * math.sin(2 * math.pi * r / num_radial_slices) + origin_point = mn.Vector3(cx, y_val, cz) + ray_set.append( + habitat_sim.geo.Ray(origin_point, center_point - origin_point) + ) + + if True: + # cast from box edges along box axes + # NOTE: using num_radial_slices per face + # for each vertical slice, select rays from each xz edge + front_to_back = mn.Vector3(aabb.right - aabb.left, 0, 0) * 1.1 + back_to_front = mn.Vector3(aabb.left - aabb.right, 0, 0) * 1.1 + left_to_right = mn.Vector3(0, 0, aabb.back - aabb.front) * 1.1 + right_to_left = mn.Vector3(0, 0, aabb.front - aabb.back) * 1.1 + for i in range(num_vertical_slices): + y_val = aabb.bottom + (i / (num_vertical_slices - 1)) * size_y + # front to back + for s in range(num_radial_slices): + z_val = aabb.back + size_z * (s / (num_radial_slices - 1)) + origin_point1 = mn.Vector3(aabb.left * 1.1, y_val, z_val) + origin_point2 = mn.Vector3(aabb.right * 1.1, y_val, z_val) + ray_set.append(habitat_sim.geo.Ray(origin_point1, front_to_back)) + ray_set.append(habitat_sim.geo.Ray(origin_point2, back_to_front)) + # left to right + for s in range(num_radial_slices): + x_val = aabb.left + size_x * (s / (num_radial_slices - 1)) + origin_point1 = mn.Vector3(x_val, y_val, aabb.front * 1.1) + origin_point2 = mn.Vector3(x_val, y_val, aabb.back * 1.1) + ray_set.append(habitat_sim.geo.Ray(origin_point1, left_to_right)) + ray_set.append(habitat_sim.geo.Ray(origin_point2, right_to_left)) + + # sphere (objects): + # TODO: compute the sphere rad size to contain the aabb: max(size)/2 + # TODO: jittered spherical sample or icosphere verts + # TODO: aim at the sphere center + + # move object to a safe (far away) location and re-orient to identity + cached_transform = obj.transformation + cached_mt = obj.motion_type + obj.motion_type = physics.MotionType.KINEMATIC + obj.translation += mn.Vector3(9000, 9000, 9000) + obj.rotation = mn.Quaternion() # identity + + # raycast: + for ray in ray_set: + # move the local ray to global space + ray.origin += obj.translation + # cast ray and get fist contact point + ray_results = sim.cast_ray(ray, max_distance=circle_rad * 2) + if ray_results.has_hits(): + surface_points.append(ray_results.hits[0].point) + ray.origin -= obj.translation + + # culling: + if cull_points: + # first compute pairwise distance + distances: List[Tuple[float, int, int]] = [] + for pix in range(len(surface_points)): + for pix2 in range(pix + 1, len(surface_points)): + # tuple (dist, index1, index2) + distances.append( + ((surface_points[pix] - surface_points[pix2]).length(), pix, pix2) + ) + remove_ixs = [] + # pairwise nearest point removal + while len(surface_points) - len(remove_ixs) > max_point_set_size: + # sort smallest distance to the top + distances.sort(key=lambda x: x[0]) + # determine which of the pair to remove by identifying the one with next closest neighbor + candidates = [distances[0][1], distances[0][2]] + remove_ix = candidates[0] + for _dist, ix1, ix2 in distances[1:]: + if ix1 in candidates: + remove_ix = ix1 + break + if ix2 in candidates: + remove_ix = ix2 + break + remove_ixs.append(remove_ix) + # remove the index from distances + distances = [tpl for tpl in distances if remove_ix not in tpl] + surface_points = [ + surface_points[ix] + for ix in range(len(surface_points)) + if ix not in remove_ixs + ] + + surface_points = obj.transform_world_pts_to_local(surface_points, link_id=-1) + + # return the object to initial state + obj.transformation = cached_transform + obj.motion_type = cached_mt + + return surface_points, ray_set + + # follow-ups: + # TODO: links + + # tests: + # TODO: scaled object + # TODO: re-oriented object + # TODO: thin structures + # TODO: L shaped couch + + +def save_interaction_points_to_markerset( + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], + interaction_points: List[mn.Vector3], +) -> None: + """ + Save the set of interaction points into the object's user_defined metadata as a MarkerSet. + """ + + obj.marker_sets.set_task_link_markerset_points( + "interaction_surface_points", "body", "primary", interaction_points + ) + + +def save_markerset_attributes( + sim: habitat_sim.Simulator, + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], +) -> None: + """ + Modify the attributes for the passed object to include the + currently edited markersets and save those attributes to disk + """ + # get the name of the attrs used to initialize the object + obj_init_attr_handle = obj.creation_attributes.handle + + if obj.is_articulated: + # save AO config + attrMgr = sim.metadata_mediator.ao_template_manager + else: + # save obj config + attrMgr = sim.metadata_mediator.object_template_manager + # get copy of initialization attributes as they were in manager, + # unmodified by scene instance values such as scale + init_attrs = attrMgr.get_template_by_handle(obj_init_attr_handle) + # TEMP TODO Remove this when fixed in Simulator + # Clean up sub-dirs being added to asset handles. + if obj.is_articulated: + init_attrs.urdf_filepath = init_attrs.urdf_filepath.split(os.sep)[-1] + init_attrs.render_asset_handle = init_attrs.render_asset_handle.split(os.sep)[ + -1 + ] + else: + init_attrs.render_asset_handle = init_attrs.render_asset_handle.split(os.sep)[ + -1 + ] + init_attrs.collision_asset_handle = init_attrs.collision_asset_handle.split( + os.sep + )[-1] + # put edited subconfig into initial attributes' markersets + markersets = init_attrs.get_marker_sets() + for subconfig_key in obj.marker_sets.get_subconfig_keys(): + markersets.save_subconfig( + subconfig_key, obj.marker_sets.get_subconfig(subconfig_key) + ) + + # reregister template + attrMgr.register_template(init_attrs, init_attrs.handle, True) + # save to original location - uses saved location in attributes + attrMgr.save_template_by_handle(init_attrs.handle, True) + + class HabitatSimInteractiveViewer(Application): # the maximum number of chars displayable in the app window # using the magnum text module. These chars are used to @@ -78,6 +285,17 @@ def __init__(self, sim_settings: Dict[str, Any]) -> None: self.contact_debug_draw = False # draw semantic region debug visualizations if present self.semantic_region_debug_draw = False + self.surface_points: List[mn.Vector3] = None + self.debug_rays: List[habitat_sim.geo.Ray] = None + self.draw_debug_rays = True + self.surface_point_obj: Union[ + physics.ManagedArticulatedObject, physics.ManagedRigidObject + ] = None + self.num_vertical_slices: int = 10 + self.num_radial_slices: int = 10 + self.previous_surface_point_compute_time = 0.0 + self.do_culling = True + self.cull_to = 20 # cache most recently loaded URDF file for quick-reload self.cached_urdf = "" @@ -188,6 +406,12 @@ def __init__(self, sim_settings: Dict[str, Any]) -> None: ): self.navmesh_config_and_recompute() + # NOTE: precompute the interaction points for all objects + # for obj in get_all_objects(self.sim): + # interaction_points, _ = find_interaction_surface_points(self.sim, obj) + # save_interaction_points_to_markerset(obj, interaction_points) + # save_markerset_attributes(self.sim, obj) + self.time_since_last_simulation = 0.0 LoggingContext.reinitialize_from_env() logger.setLevel("INFO") @@ -262,6 +486,52 @@ def debug_draw(self): ) self.draw_region_debug(debug_line_render) + # if self.surface_points is not None and len(self.surface_points) > 1: + # centroid = mn.Vector3() + # for point in self.surface_points: + # centroid += point + # centroid /= len(self.surface_points) + # debug_line_render.push_transform(self.surface_point_obj.transformation) + # for point in self.surface_points: + # debug_line_render.draw_circle( + # translation=point, + # radius=0.005, + # color=mn.Color4.yellow(), + # normal=centroid - point, + # ) + # if self.draw_debug_rays: + # for ray in self.debug_rays: + # debug_line_render.draw_transformed_line( + # ray.origin, + # ray.origin + ray.direction, + # mn.Color4.green(), + # ) + # debug_line_render.pop_transform() + # draw any active marker_sets with interaction_surface_points + for obj in get_all_objects(self.sim): + if obj.marker_sets.has_taskset("interaction_surface_points"): + points = obj.marker_sets.get_task_link_markerset_points( + "interaction_surface_points", "body", "primary" + ) + global_points = obj.transform_local_pts_to_world(points, link_id=-1) + # global_points = [obj.transformation.transform_point(point) for point in points] + # dif = [pix for pix in range(len(global_points)) if global_points[pix] != alt_global_points[pix]] + # if len(dif) > 0: + # breakpoint() + centroid = mn.Vector3() + for point in global_points: + centroid += point + centroid /= len(points) + # debug_line_render.push_transform(obj.transformation) + for point in global_points: + debug_line_render.draw_circle( + translation=point, + radius=0.005, + color=mn.Color4.blue(), + normal=centroid - point, + ) + # debug_line_render.pop_transform() + def draw_event( self, simulation_call: Optional[Callable] = None, @@ -558,72 +828,86 @@ def key_press_event(self, event: Application.KeyEvent) -> None: logger.info(f"Command: toggle Bullet debug draw: {self.debug_bullet_draw}") elif key == pressed.C: - if shift_pressed: - self.contact_debug_draw = not self.contact_debug_draw - logger.info( - f"Command: toggle contact debug draw: {self.contact_debug_draw}" - ) + if alt_pressed: + self.do_culling = not self.do_culling + print(f"do_culling = {self.do_culling}") + elif shift_pressed: + self.cull_to -= 1 else: - # perform a discrete collision detection pass and enable contact debug drawing to visualize the results - logger.info( - "Command: perform discrete collision detection and visualize active contacts." + self.cull_to += 1 + self.cull_to = max(3, self.cull_to) + if self.surface_point_obj is not None: + start_time = time.time() + self.surface_points, self.debug_rays = find_interaction_surface_points( + self.sim, + self.surface_point_obj, + num_radial_slices=self.num_radial_slices, + num_vertical_slices=self.num_vertical_slices, + cull_points=self.do_culling, + max_point_set_size=self.cull_to, ) - self.sim.perform_discrete_collision_detection() - self.contact_debug_draw = True - # TODO: add a nice log message with concise contact pair naming. + self.previous_surface_point_compute_time = time.time() - start_time elif key == pressed.T: - # load URDF - fixed_base = alt_pressed - urdf_file_path = "" - if shift_pressed and self.cached_urdf: - urdf_file_path = self.cached_urdf - else: - urdf_file_path = input("Load URDF: provide a URDF filepath:").strip() - - if not urdf_file_path: - logger.warn("Load URDF: no input provided. Aborting.") - elif not urdf_file_path.endswith((".URDF", ".urdf")): - logger.warn("Load URDF: input is not a URDF. Aborting.") - elif os.path.exists(urdf_file_path): - self.cached_urdf = urdf_file_path - aom = self.sim.get_articulated_object_manager() - ao = aom.add_articulated_object_from_urdf( - urdf_file_path, - fixed_base, - 1.0, - 1.0, - True, - maintain_link_order=False, - intertia_from_urdf=False, + # toggle debug ray display + self.draw_debug_rays = not self.draw_debug_rays + print(f"draw_debug_rays = {self.draw_debug_rays}") + + elif key == pressed.F: + # save the current interaction_surface_points as a MarkerSet + if self.surface_point_obj is not None and self.surface_points is not None: + save_interaction_points_to_markerset( + self.surface_point_obj, self.surface_points ) - ao.translation = ( - self.default_agent.scene_node.transformation.transform_point( - [0.0, 1.0, -1.5] - ) + print( + f"Saved {len(self.surface_points)} surface points to marker set for objects {self.surface_point_obj.handle}" ) - # check removal and auto-creation - joint_motor_settings = habitat_sim.physics.JointMotorSettings( - position_target=0.0, - position_gain=1.0, - velocity_target=0.0, - velocity_gain=1.0, - max_impulse=1000.0, - ) - existing_motor_ids = ao.existing_joint_motor_ids - for motor_id in existing_motor_ids: - ao.remove_joint_motor(motor_id) - ao.create_all_motors(joint_motor_settings) - else: - logger.warn("Load URDF: input file not found. Aborting.") + if shift_pressed: + save_markerset_attributes(self.sim, self.surface_point_obj) + print("Also saved markerset to config.") elif key == pressed.M: self.cycle_mouse_mode() logger.info(f"Command: mouse mode set to {self.mouse_interaction}") elif key == pressed.V: - self.invert_gravity() - logger.info("Command: gravity inverted") + # increment vertical slice count + if shift_pressed: + self.num_vertical_slices -= 1 + else: + self.num_vertical_slices += 1 + self.num_vertical_slices = max(1, self.num_vertical_slices) + # recompute surface points + if self.surface_point_obj is not None: + start_time = time.time() + self.surface_points, self.debug_rays = find_interaction_surface_points( + self.sim, + self.surface_point_obj, + num_radial_slices=self.num_radial_slices, + num_vertical_slices=self.num_vertical_slices, + cull_points=self.do_culling, + max_point_set_size=self.cull_to, + ) + self.previous_surface_point_compute_time = time.time() - start_time + elif key == pressed.R: + # increment radial slice count + if shift_pressed: + self.num_radial_slices -= 1 + else: + self.num_radial_slices += 1 + self.num_radial_slices = max(3, self.num_radial_slices) + # recompute surface points + if self.surface_point_obj is not None: + start_time = time.time() + self.surface_points, self.debug_rays = find_interaction_surface_points( + self.sim, + self.surface_point_obj, + num_radial_slices=self.num_radial_slices, + num_vertical_slices=self.num_vertical_slices, + cull_points=self.do_culling, + max_point_set_size=self.cull_to, + ) + self.previous_surface_point_compute_time = time.time() - start_time elif key == pressed.N: # (default) - toggle navmesh visualization # NOTE: (+ALT) - re-sample the agent position on the NavMesh @@ -714,7 +998,7 @@ def mouse_press_event(self, event: Application.MouseEvent) -> None: physics_enabled = self.sim.get_physics_simulation_library() # if interactive mode is True -> GRAB MODE - if self.mouse_interaction == MouseMode.GRAB and physics_enabled: + if physics_enabled: render_camera = self.render_camera.render_camera ray = render_camera.unproject(self.get_mouse_position(event.position)) raycast_results = self.sim.cast_ray(ray=ray) @@ -765,33 +1049,54 @@ def mouse_press_event(self, event: Application.MouseEvent) -> None: # done checking for AO if hit_object >= 0: - node = self.default_agent.scene_node - constraint_settings = physics.RigidConstraintSettings() - - constraint_settings.object_id_a = hit_object - constraint_settings.link_id_a = ao_link - constraint_settings.pivot_a = object_pivot - constraint_settings.frame_a = ( - object_frame.to_matrix() @ node.rotation.to_matrix() - ) - constraint_settings.frame_b = node.rotation.to_matrix() - constraint_settings.pivot_b = hit_info.point - - # by default use a point 2 point constraint - if event.button == button.RIGHT: - constraint_settings.constraint_type = ( - physics.RigidConstraintType.Fixed + if self.mouse_interaction == MouseMode.GRAB: + node = self.default_agent.scene_node + constraint_settings = physics.RigidConstraintSettings() + + constraint_settings.object_id_a = hit_object + constraint_settings.link_id_a = ao_link + constraint_settings.pivot_a = object_pivot + constraint_settings.frame_a = ( + object_frame.to_matrix() @ node.rotation.to_matrix() ) + constraint_settings.frame_b = node.rotation.to_matrix() + constraint_settings.pivot_b = hit_info.point - grip_depth = ( - hit_info.point - render_camera.node.absolute_translation - ).length() + # by default use a point 2 point constraint + if event.button == button.RIGHT: + constraint_settings.constraint_type = ( + physics.RigidConstraintType.Fixed + ) - self.mouse_grabber = MouseGrabber( - constraint_settings, - grip_depth, - self.sim, - ) + grip_depth = ( + hit_info.point - render_camera.node.absolute_translation + ).length() + + self.mouse_grabber = MouseGrabber( + constraint_settings, + grip_depth, + self.sim, + ) + elif event.button == button.RIGHT: + # right click in LOOK + self.surface_point_obj = get_obj_from_id( + self.sim, hit_object + ) + start_time = time.time() + ( + self.surface_points, + self.debug_rays, + ) = find_interaction_surface_points( + self.sim, + self.surface_point_obj, + num_radial_slices=self.num_radial_slices, + num_vertical_slices=self.num_vertical_slices, + cull_points=self.do_culling, + max_point_set_size=self.cull_to, + ) + self.previous_surface_point_compute_time = ( + time.time() - start_time + ) else: logger.warn("Oops, couldn't find the hit object. That's odd.") # end if didn't hit the scene @@ -950,6 +1255,11 @@ def draw_text(self, sensor_spec): Sensor Type: {sensor_type_string} Sensor Subtype: {sensor_subtype_string} Mouse Interaction Mode: {mouse_mode_string} +Radial Slices: {self.num_radial_slices} +Vertical Slices: {self.num_vertical_slices} +Prev Compute Time: {self.previous_surface_point_compute_time} +Culling to: {self.cull_to if self.do_culling else "NA"} +Num Points: {len(self.surface_points) if self.surface_points is not None else None} """ ) self.shader.draw(self.window_text.mesh) diff --git a/src/esp/bindings/PhysicsObjectBindings.cpp b/src/esp/bindings/PhysicsObjectBindings.cpp index eb046db50b..a5ea15f223 100644 --- a/src/esp/bindings/PhysicsObjectBindings.cpp +++ b/src/esp/bindings/PhysicsObjectBindings.cpp @@ -78,7 +78,7 @@ void declareBasePhysicsObjectWrapper(py::module& m, &PhysObjWrapper::transformLocalPointsToWorld, R"(Given the list of passed points in this object's local space, return those points transformed to world space. The link_id is for articulated - objects and is ignored for rigid objects and stages )", + objects and is ignored for rigid objects and stages. link_id==-1 uses the AO's base transform.)", "ls_points"_a, "link_id"_a) .def_property( "rotation", &PhysObjWrapper::getRotation, diff --git a/src/esp/physics/ArticulatedObject.h b/src/esp/physics/ArticulatedObject.h index d3b63438dd..1753a38c2d 100644 --- a/src/esp/physics/ArticulatedObject.h +++ b/src/esp/physics/ArticulatedObject.h @@ -493,7 +493,10 @@ class ArticulatedObject : public esp::physics::PhysicsObjectBase { */ std::vector transformLocalPointsToWorld( const std::vector& points, - int linkId) const override { + int linkId = -1) const override { + if (linkId == -1) { + return this->baseLink_->transformLocalPointsToWorld(points, -1); + } auto linkIter = links_.find(linkId); ESP_CHECK(linkIter != links_.end(), "ArticulatedObject::getLinkVisualSceneNodes - no link found with " @@ -511,7 +514,10 @@ class ArticulatedObject : public esp::physics::PhysicsObjectBase { */ std::vector transformWorldPointsToLocal( const std::vector& points, - int linkId) const override { + int linkId = -1) const override { + if (linkId == -1) { + return this->baseLink_->transformWorldPointsToLocal(points, -1); + } auto linkIter = links_.find(linkId); ESP_CHECK(linkIter != links_.end(), "ArticulatedObject::getLinkVisualSceneNodes - no link found with " @@ -549,19 +555,25 @@ class ArticulatedObject : public esp::physics::PhysicsObjectBase { for (const auto& linkEntry : taskEntry.second) { const std::string linkName = linkEntry.first; int linkId = getLinkIdFromName(linkName); - auto linkIter = links_.find(linkId); - ESP_CHECK( - linkIter != links_.end(), - "ArticulatedObject::getMarkerPointsGlobal - no link found with " - "linkId =" - << linkId); + // locally access the unique pointer's payload + const esp::physics::ArticulatedLink* aoLink; + if (linkId == -1) { + aoLink = baseLink_.get(); + } else { + auto linkIter = links_.find(linkId); + ESP_CHECK( + linkIter != links_.end(), + "ArticulatedObject::getMarkerPointsGlobal - no link found with " + "linkId =" + << linkId); + aoLink = linkIter->second.get(); + } std::unordered_map> perLinkMap; // for each set in link for (const auto& markersEntry : linkEntry.second) { const std::string markersName = markersEntry.first; perLinkMap[markersName] = - linkIter->second->transformLocalPointsToWorld(markersEntry.second, - linkId); + aoLink->transformLocalPointsToWorld(markersEntry.second, linkId); } perTaskMap[linkName] = perLinkMap; } diff --git a/src/esp/physics/PhysicsObjectBase.h b/src/esp/physics/PhysicsObjectBase.h index fa813a862b..ade6e34b3c 100644 --- a/src/esp/physics/PhysicsObjectBase.h +++ b/src/esp/physics/PhysicsObjectBase.h @@ -229,11 +229,11 @@ class PhysicsObjectBase : public Magnum::SceneGraph::AbstractFeature3D { */ virtual std::vector transformLocalPointsToWorld( const std::vector& points, - CORRADE_UNUSED int linkID) const { + CORRADE_UNUSED int linkID = -1) const { std::vector wsPoints; wsPoints.reserve(points.size()); Mn::Vector3 objScale = getScale(); - Mn::Matrix4 worldTransform = getTransformation(); + Mn::Matrix4 worldTransform = node().absoluteTransformation(); for (const auto& lsPoint : points) { wsPoints.emplace_back(worldTransform.transformPoint(lsPoint * objScale)); } @@ -249,11 +249,11 @@ class PhysicsObjectBase : public Magnum::SceneGraph::AbstractFeature3D { */ virtual std::vector transformWorldPointsToLocal( const std::vector& points, - CORRADE_UNUSED int linkID) const { + CORRADE_UNUSED int linkID = -1) const { std::vector lsPoints; lsPoints.reserve(points.size()); Mn::Vector3 objScale = getScale(); - Mn::Matrix4 worldTransform = getTransformation(); + Mn::Matrix4 worldTransform = node().absoluteTransformation(); for (const auto& wsPoint : points) { lsPoints.emplace_back(worldTransform.inverted().transformPoint(wsPoint) / objScale); diff --git a/tools/compute_interaction_points.py b/tools/compute_interaction_points.py new file mode 100644 index 0000000000..961bffa1a4 --- /dev/null +++ b/tools/compute_interaction_points.py @@ -0,0 +1,279 @@ +import math +import os +from typing import Any, Dict, List, Tuple, Union + +# NOTE: (requires habitat-lab) get metadata for semantics +import habitat.sims.habitat_simulator.sim_utilities as sutils +import magnum as mn + +import habitat_sim +from habitat_sim import Simulator, physics +from habitat_sim.metadata import MetadataMediator +from habitat_sim.utils.settings import default_sim_settings, make_cfg + + +def find_interaction_surface_points( + sim: habitat_sim.Simulator, + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], + num_vertical_slices: int = 10, + num_radial_slices: int = 10, + cull_points=True, + max_point_set_size: int = 20, +) -> Tuple[List[mn.Vector3], List[mn.Vector3]]: + """ + Use raycasting to find a set of points on the lateral surfaces of an object. + + :return: the list of interaction points and the list of cast rays + """ + + assert num_vertical_slices >= 1, "Must at least slice in half." + assert num_radial_slices >= 4, "Must at least form a 2D simplex." + assert max_point_set_size >= 3, "Must at least form a 2D simplex." + + surface_points: List[mn.Vector3] = [] + + # compute the ray set: + aabb = obj.aabb + ray_set: List[habitat_sim.geo.Ray] = [] + + # compute the circle size to contain the aabb: xz diagonal length+10%. + # Used for max ray distance and cylinder sampling. + size_x = aabb.size_x() + size_y = aabb.size_y() + size_z = aabb.size_z() + circle_rad = math.sqrt((size_x / 2.0) ** 2 + (size_z / 2.0) ** 2) * 1.1 + if False: + # cylinder (furniture): + # for each vertical slice, select rays in a circle + for i in range(1, num_vertical_slices + 1): + y_val = aabb.bottom + (i / num_vertical_slices) * size_y + center_point = mn.Vector3(0, y_val, 0) + for r in range(num_radial_slices): + cx = circle_rad * math.cos(2 * math.pi * r / num_radial_slices) + cz = circle_rad * math.sin(2 * math.pi * r / num_radial_slices) + origin_point = mn.Vector3(cx, y_val, cz) + ray_set.append( + habitat_sim.geo.Ray(origin_point, center_point - origin_point) + ) + + if True: + # cast from box edges along box axes + # NOTE: using num_radial_slices per face + # for each vertical slice, select rays from each xz edge + front_to_back = mn.Vector3(aabb.right - aabb.left, 0, 0) * 1.1 + back_to_front = mn.Vector3(aabb.left - aabb.right, 0, 0) * 1.1 + left_to_right = mn.Vector3(0, 0, aabb.back - aabb.front) * 1.1 + right_to_left = mn.Vector3(0, 0, aabb.front - aabb.back) * 1.1 + for i in range(num_vertical_slices): + y_val = aabb.bottom + (i / (num_vertical_slices - 1)) * size_y + # front to back + for s in range(num_radial_slices): + z_val = aabb.back + size_z * (s / (num_radial_slices - 1)) + origin_point1 = mn.Vector3(aabb.left * 1.1, y_val, z_val) + origin_point2 = mn.Vector3(aabb.right * 1.1, y_val, z_val) + ray_set.append(habitat_sim.geo.Ray(origin_point1, front_to_back)) + ray_set.append(habitat_sim.geo.Ray(origin_point2, back_to_front)) + # left to right + for s in range(num_radial_slices): + x_val = aabb.left + size_x * (s / (num_radial_slices - 1)) + origin_point1 = mn.Vector3(x_val, y_val, aabb.front * 1.1) + origin_point2 = mn.Vector3(x_val, y_val, aabb.back * 1.1) + ray_set.append(habitat_sim.geo.Ray(origin_point1, left_to_right)) + ray_set.append(habitat_sim.geo.Ray(origin_point2, right_to_left)) + + # sphere (objects): + # TODO: compute the sphere rad size to contain the aabb: max(size)/2 + # TODO: jittered spherical sample or icosphere verts + # TODO: aim at the sphere center + + # move object to a safe (far away) location and re-orient to identity + cached_transform = obj.transformation + cached_mt = obj.motion_type + obj.motion_type = physics.MotionType.KINEMATIC + obj.translation += mn.Vector3(9000, 9000, 9000) + obj.rotation = mn.Quaternion() # identity + + # raycast: + for ray in ray_set: + # move the local ray to global space + ray.origin += obj.translation + # cast ray and get fist contact point + ray_results = sim.cast_ray(ray, max_distance=circle_rad * 2) + if ray_results.has_hits(): + surface_points.append(ray_results.hits[0].point) + ray.origin -= obj.translation + + # culling: + if cull_points: + # first compute pairwise distance + distances: List[Tuple[float, int, int]] = [] + for pix in range(len(surface_points)): + for pix2 in range(pix + 1, len(surface_points)): + # tuple (dist, index1, index2) + distances.append( + ((surface_points[pix] - surface_points[pix2]).length(), pix, pix2) + ) + remove_ixs = [] + # pairwise nearest point removal + while len(surface_points) - len(remove_ixs) > max_point_set_size: + # sort smallest distance to the top + distances.sort(key=lambda x: x[0]) + # determine which of the pair to remove by identifying the one with next closest neighbor + candidates = [distances[0][1], distances[0][2]] + remove_ix = candidates[0] + for _dist, ix1, ix2 in distances[1:]: + if ix1 in candidates: + remove_ix = ix1 + break + if ix2 in candidates: + remove_ix = ix2 + break + remove_ixs.append(remove_ix) + # remove the index from distances + distances = [tpl for tpl in distances if remove_ix not in tpl] + surface_points = [ + surface_points[ix] + for ix in range(len(surface_points)) + if ix not in remove_ixs + ] + + surface_points = obj.transform_world_pts_to_local(surface_points, link_id=-1) + + # return the object to initial state + obj.transformation = cached_transform + obj.motion_type = cached_mt + + return surface_points, ray_set + + # follow-ups: + # TODO: links + + # tests: + # TODO: scaled object + # TODO: re-oriented object + # TODO: thin structures + # TODO: L shaped couch + + +def save_interaction_points_to_markerset( + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], + interaction_points: List[mn.Vector3], +) -> None: + """ + Save the set of interaction points into the object's user_defined metadata as a MarkerSet. + """ + + obj.marker_sets.set_task_link_markerset_points( + "interaction_surface_points", "body", "primary", interaction_points + ) + + +def save_markerset_attributes( + sim: habitat_sim.Simulator, + obj: Union[physics.ManagedRigidObject, physics.ManagedArticulatedObject], +) -> None: + """ + Modify the attributes for the passed object to include the + currently edited markersets and save those attributes to disk + """ + # get the name of the attrs used to initialize the object + obj_init_attr_handle = obj.creation_attributes.handle + + if obj.is_articulated: + # save AO config + attrMgr = sim.metadata_mediator.ao_template_manager + else: + # save obj config + attrMgr = sim.metadata_mediator.object_template_manager + # get copy of initialization attributes as they were in manager, + # unmodified by scene instance values such as scale + init_attrs = attrMgr.get_template_by_handle(obj_init_attr_handle) + # TEMP TODO Remove this when fixed in Simulator + # Clean up sub-dirs being added to asset handles. + if obj.is_articulated: + init_attrs.urdf_filepath = init_attrs.urdf_filepath.split(os.sep)[-1] + init_attrs.render_asset_handle = init_attrs.render_asset_handle.split(os.sep)[ + -1 + ] + else: + init_attrs.render_asset_handle = init_attrs.render_asset_handle.split(os.sep)[ + -1 + ] + init_attrs.collision_asset_handle = init_attrs.collision_asset_handle.split( + os.sep + )[-1] + # put edited subconfig into initial attributes' markersets + markersets = init_attrs.get_marker_sets() + for subconfig_key in obj.marker_sets.get_subconfig_keys(): + markersets.save_subconfig( + subconfig_key, obj.marker_sets.get_subconfig(subconfig_key) + ) + + # reregister template + attrMgr.register_template(init_attrs, init_attrs.handle, True) + # save to original location - uses saved location in attributes + attrMgr.save_template_by_handle(init_attrs.handle, True) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + + parser.add_argument( + "--dataset", + default="default", + type=str, + metavar="DATASET", + help='dataset configuration file to use (default: "default")', + ) + + parser.add_argument( + "--scenes", + nargs="+", + type=str, + help="A subset of scene names to process. Limits the iteration to less than the full set of scenes.", + default=None, + ) + + args = parser.parse_args() + + # create an initial simulator config + sim_settings: Dict[str, Any] = default_sim_settings + sim_settings["scene_dataset_config_file"] = args.dataset + cfg = make_cfg(sim_settings) + + # pre-initialize a MetadataMediator to iterate over scenes + mm = MetadataMediator() + mm.active_dataset = args.dataset + + target_scenes = mm.get_scene_handles() + if args.scenes is not None: + target_scenes = args.scenes + num_scenes = len(target_scenes) + + for s_ix, scene_handle in enumerate(target_scenes): + print("=================================================================") + print( + f"Setting up scene for {scene_handle} ({s_ix}|{num_scenes} = {s_ix/float(num_scenes)*100}%)" + ) + + cfg.sim_cfg.scene_id = scene_handle + print(" - init") + with Simulator(cfg) as sim: + objects = sutils.get_all_objects(sim) + print(f" - processing {len(objects)} objects:") + for oix, obj in enumerate(objects): + print(f" - obj ({oix}/{len(objects)}) {obj.handle}") + # if not obj.marker_sets.has_taskset("interaction_surface_points"): + print(" - computing interaction points") + surface_points, debug_rays = find_interaction_surface_points( + sim, + obj, + num_radial_slices=10, + num_vertical_slices=10, + cull_points=True, + max_point_set_size=20, + ) + save_interaction_points_to_markerset(obj, surface_points) + save_markerset_attributes(sim, obj)