Skip to content

visualize_object

main(random_selection=False, headless=False, short_exec=False)

Visualizes object as specified by its USD path, @usd_path. If None if specified, will instead result in an object selection from OmniGibson's object dataset

Source code in examples/objects/visualize_object.py
def main(random_selection=False, headless=False, short_exec=False):
    """
    Visualizes object as specified by its USD path, @usd_path. If None if specified, will instead
    result in an object selection from OmniGibson's object dataset
    """
    logging.info("*" * 80 + "\nDescription:" + main.__doc__ + "*" * 80)

    # Assuming that if random_selection=True, headless=True, short_exec=True, we are calling it from tests and we
    # do not want to parse args (it would fail because the calling function is pytest "testfile.py")
    usd_path = None
    if not (random_selection and headless and short_exec):
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--usd_path",
            default=None,
            help="USD Model to load",
        )
        args = parser.parse_args()
        usd_path = args.usd_path

    # Define objects to load
    light0_cfg = OrderedDict(
        type="LightObject",
        light_type="Sphere",
        name="sphere_light0",
        radius=0.01,
        intensity=1e5,
        position=[-2.0, -2.0, 2.0],
    )

    light1_cfg = OrderedDict(
        type="LightObject",
        light_type="Sphere",
        name="sphere_light1",
        radius=0.01,
        intensity=1e5,
        position=[-2.0, 2.0, 2.0],
    )

    # Make sure we have a valid usd path
    if usd_path is None:
        # Select a category to load
        available_obj_categories = get_all_object_categories()
        obj_category = choose_from_options(options=available_obj_categories, name="object category",
                                           random_selection=random_selection)

        # Select a model to load
        available_obj_models = get_object_models_of_category(obj_category)
        obj_model = choose_from_options(options=available_obj_models, name="object model",
                                        random_selection=random_selection)

        kwargs = {
            "type": "DatasetObject",
            "category": obj_category,
            "model": obj_model,
        }
    else:
        kwargs = {
            "type": "USDObject",
            "usd_path": usd_path,
        }

    # Import the desired object
    obj_cfg = OrderedDict(
        **kwargs,
        name="obj",
        usd_path=usd_path,
        visual_only=True,
        position=[0, 0, 10.0],
    )

    # Create the scene config to load -- empty scene
    cfg = {
        "scene": {
            "type": "Scene",
        },
        "objects": [light0_cfg, light1_cfg, obj_cfg],
    }

    # Create the environment
    env = og.Environment(configs=cfg)

    # Set camera to appropriate viewing pose
    og.sim.viewer_camera.set_position_orientation(
        position=np.array([-0.00913503, -1.95750906,  1.36407314]),
        orientation=np.array([0.6350064 , 0.        , 0.        , 0.77250687]),
    )

    # Grab the object references
    obj = env.scene.object_registry("name", "obj")

    # Standardize the scale of the object so it fits in a [1,1,1] box -- note that we have to stop the simulator
    # in order to set the scale
    extents = obj.aabb_extent
    og.sim.stop()
    obj.scale = (np.ones(3) / extents).min()
    og.sim.play()
    env.step(np.array([]))

    # Move the object so that its center is at [0, 0, 1]
    center_offset = obj.aabb_center - obj.get_position() + np.array([0, 0, 1.0])
    obj.set_position(center_offset)

    # Allow the user to easily move the camera around
    og.sim.enable_viewer_camera_teleoperation()

    # Rotate the object in place
    steps_per_rotate = 360
    steps_per_joint = steps_per_rotate / 10
    max_steps = 100 if short_exec else 10000
    for i in range(max_steps):
        z_angle = (2 * np.pi * (i % steps_per_rotate) / steps_per_rotate)
        quat = T.euler2quat(np.array([0, 0, z_angle]))
        pos = T.quat2mat(quat) @ center_offset
        if obj.n_dof > 0:
            frac = (i % steps_per_joint) / steps_per_joint
            j_frac = -1.0 + 2.0 * frac if (i // steps_per_joint) % 2 == 0 else 1.0 - 2.0 * frac
            obj.set_joint_positions(positions=j_frac * np.ones(obj.n_dof), normalized=True, target=False)
            obj.keep_still()
        obj.set_position_orientation(position=pos, orientation=quat)
        env.step(np.array([]))