HITL - Add picture-in-picture viewports. (facebookresearch#1958)

* Add viewports.

* Add viewports to rearrange_v2.

* Improve viewport rect description.

examples/hitl/rearrange_v2/rearrange_v2.py

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 
+from __future__ import annotations
+
 # Must call this before importing Habitat or Magnum.
 # fmt: off
 import ctypes
@@ -13,7 +15,6 @@
 sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)
 # fmt: on
 
-
 from typing import List, Optional
 
 # This registers collaboration episodes into this application.
@@ -45,6 +46,7 @@
 from habitat_sim.utils.common import quat_from_magnum, quat_to_coeffs
 
 UP = mn.Vector3(0, 1, 0)
+PIP_VIEWPORT_ID = 0  # ID of the picture-in-picture viewport that shows other agent's perspective.
 
 
 class DataLogger:
@@ -127,6 +129,7 @@ def __init__(
         self.show_gui_text = True
         self.task_instruction = ""
         self.signal_change_episode = False
+        self.pip_initialized = False
 
         # If in remote mode, get the remote input. Else get the server (local) input.
         self.gui_input = (
@@ -186,6 +189,28 @@ def update(self, dt: float):
         self.ui.update()
         self.ui.draw_ui()
 
+    def draw_pip_viewport(self, pip_user_data: UserData):
+        """
+        Draw a picture-in-picture viewport showing another agent's perspective.
+        """
+        # Lazy init:
+        if not self.pip_initialized:
+            self.pip_initialized = True
+
+            # Define picture-in-picture (PIP) viewport.
+            self.app_service.client_message_manager.set_viewport_properties(
+                viewport_id=PIP_VIEWPORT_ID,
+                viewport_rect_xywh=[0.8, 0.02, 0.18, 0.18],
+                destination_mask=Mask.from_index(self.user_index),
+            )
+
+        # Show picture-in-picture (PIP) viewport.
+        self.app_service.client_message_manager.show_viewport(
+            viewport_id=PIP_VIEWPORT_ID,
+            cam_transform=pip_user_data.cam_transform,
+            destination_mask=Mask.from_index(self.user_index),
+        )
+
     def _get_camera_lookat_pos(self) -> mn.Vector3:
         agent_root = get_agent_art_obj_transform(
             self.app_service.sim,
@@ -396,6 +421,11 @@ def sim_update(self, dt: float, post_sim_update_dict):
 
         self._sps_tracker.increment()
 
+        # Draw the picture-in-picture showing other agent's perspective.
+        if self._users.max_user_count == 2:
+            self._user_data[0].draw_pip_viewport(self._user_data[1])
+            self._user_data[1].draw_pip_viewport(self._user_data[0])
+
         if not self._paused:
             for user_index in self._users.indices(Mask.ALL):
                 self._user_data[user_index].update(dt)

habitat-hitl/habitat_hitl/core/client_message_manager.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
-from typing import Final, List, Optional, Union
+from typing import Any, Dict, Final, List, Optional, Union
 
 import magnum as mn
 
@@ -242,6 +242,50 @@ def set_server_keyframe_id(
             message = self._messages[user_index]
             message["serverKeyframeId"] = keyframe_id
 
+    def set_viewport_properties(
+        self,
+        viewport_id: int,
+        viewport_rect_xywh: List[float],
+        destination_mask: Mask = Mask.ALL,
+    ):
+        r"""
+        Set the properties of a viewport. Unlike show_viewport(), this does not have to be called every frame.
+        Use viewport_id '-1' to edit the default viewport.
+
+        viewport_id: Unique identifier of the viewport.
+        viewport_rect_xywh: Viewport rect (x position, y position, width, height).
+                            In window normalized coordinates, i.e. all values in range [0,1] relative to window size.
+        """
+        for user_index in self._users.indices(destination_mask):
+            message = self._messages[user_index]
+            viewport_properties = _obtain_viewport_properties(
+                message, viewport_id
+            )
+            viewport_properties["rect"] = viewport_rect_xywh
+
+    def show_viewport(
+        self,
+        viewport_id: int,
+        cam_transform: mn.Matrix4,
+        destination_mask: Mask = Mask.ALL,
+    ):
+        """
+        Show a picture-in-picture viewport rendering the specified camera matrix.
+        This must be repeatedly called for the viewport to stay visible.
+        The viewport_id '-1' is reserved for the main viewport. It is always visible.
+        Use set_viewport_properties() to configure the viewport.
+        """
+        assert viewport_id != -1
+        for user_index in self._users.indices(destination_mask):
+            message = self._messages[user_index]
+            viewport_properties = _obtain_viewport_properties(
+                message, viewport_id
+            )
+            viewport_properties["enabled"] = True
+            viewport_properties["camera"] = _create_transform_dict(
+                cam_transform
+            )
+
     def update_navmesh_triangles(
         self,
         triangle_vertices: List[List[float]],
@@ -290,3 +334,25 @@ def update_camera_transform(
                 rot[2],
                 rot[3],
             ]
+
+
+def _create_transform_dict(transform: mn.Matrix4) -> Dict[str, List[float]]:
+    """Create a message dictionary from a transform."""
+    p = transform.translation
+    r = mn.Quaternion.from_matrix(transform.rotation())
+    rv = r.vector
+    return {
+        "translation": [p[0], p[1], p[2]],
+        "rotation": [r.scalar, rv[0], rv[1], rv[2]],
+    }
+
+
+def _obtain_viewport_properties(
+    message: Message, viewport_id: int
+) -> Dict[str, Any]:
+    """Get or create the properties dict of an object_id."""
+    if "viewports" not in message:
+        message["viewports"] = {}
+    if viewport_id not in message["viewports"]:
+        message["viewports"][viewport_id] = {}
+    return message["viewports"][viewport_id]