From 618373a1dd4c11428717a7d6a98e637ff39da657 Mon Sep 17 00:00:00 2001
From: Kavit Shah <kavitsha@meta.com>
Date: Mon, 2 Oct 2023 12:23:51 -0700
Subject: [PATCH] Enable passing and receiving subplots for viz

---
 perception/fairotag/python/fairotag/camera.py |  2 +-
 perception/fairotag/python/fairotag/scene.py  | 10 +++++----
 perception/fairotag/python/fairotag/viz.py    | 22 +++++++++++++------
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/perception/fairotag/python/fairotag/camera.py b/perception/fairotag/python/fairotag/camera.py
index 12f91c550e..9d5349f083 100644
--- a/perception/fairotag/python/fairotag/camera.py
+++ b/perception/fairotag/python/fairotag/camera.py
@@ -126,7 +126,7 @@ def render_markers(self, img, markers=None):
                 tvec = m.pose.translation()[None, :]
                 rvec = m.pose.so3().log()[None, :]
                 length = m.length / 2.0
-                img_rend = cv2.aruco.drawAxis(
+                img_rend = cv2.drawFrameAxes(
                     img_rend,
                     self._intrinsics2matrix(self.intrinsics),
                     self.intrinsics.coeffs,
diff --git a/perception/fairotag/python/fairotag/scene.py b/perception/fairotag/python/fairotag/scene.py
index 4c9b66653e..0590a629ac 100644
--- a/perception/fairotag/python/fairotag/scene.py
+++ b/perception/fairotag/python/fairotag/scene.py
@@ -31,7 +31,7 @@ class Object:
     pose_in_frame: sp.SE3
     size: float
     is_anchor: bool = False
-    is_visible: bool = False
+    is_visible: bool = True
 
 
 @dataclass
@@ -64,6 +64,7 @@ def __init__(self, camera_noise=None, calib_noise=None):
 
         # Default world frame
         f0 = Frame("world", sp.SE3())
+        f0.is_visible = True
         self._frames["world"] = f0
 
     # Scene construction
@@ -344,8 +345,8 @@ def calibrate_extrinsics(
             self.clear_snapshots()
 
     # Rendering
-    def visualize(self, show_marker_id=False):
-        viz = SceneViz()
+    def visualize(self, fig=None, show_marker_id=False, should_return=False):
+        viz = SceneViz(fig=fig)
 
         # Draw markers & cameras
         for obj_name, obj in self._objects.items():
@@ -358,6 +359,7 @@ def visualize(self, show_marker_id=False):
             else:
                 viz.draw_camera(obj.pose, size=obj.size)
 
+
         # Draw frames
         for frame_name, frame in self._frames.items():
             if not frame.is_visible:
@@ -365,7 +367,7 @@ def visualize(self, show_marker_id=False):
             viz.draw_axes(frame.pose)
 
         # Show
-        viz.show()
+        return viz.show(should_return=should_return)
 
     # Save / Load TODO
     def __getstate__(self):
diff --git a/perception/fairotag/python/fairotag/viz.py b/perception/fairotag/python/fairotag/viz.py
index 74e91be056..6c1ad1955d 100644
--- a/perception/fairotag/python/fairotag/viz.py
+++ b/perception/fairotag/python/fairotag/viz.py
@@ -2,13 +2,17 @@
 import sophus as sp
 import matplotlib.pyplot as plt
 
-DEFAULT_AXIS_LENGTH = 0.1
+DEFAULT_AXIS_LENGTH = 0.3
 
 
 class SceneViz:
-    def __init__(self):
-        self.fig = plt.figure()
-        self.ax = plt.axes(projection="3d")
+    def __init__(self, fig=None, plot_config=112):
+        self.fig=None
+        if fig:
+            self.fig = fig
+        else:
+            self.fig = plt.figure()
+        self.ax = self.fig.add_subplot(1,2,2,projection="3d")
         self.max = np.zeros(3)
         self.min = np.zeros(3)
 
@@ -80,7 +84,7 @@ def draw_marker(self, pose, id, length, color="k", show_id=False):
             pos = pose.translation()
             self.ax.text(pos[0], pos[1], pos[2], id, color="b")
 
-    def show(self):
+    def show(self, should_return=False):
         # Set limits
         mid = (self.max + self.min) / 2.0
         r = max(np.max(self.max - mid), np.max(mid - self.min))
@@ -88,5 +92,9 @@ def show(self):
         self.ax.set_ylim(mid[1] - r, mid[1] + r)
         self.ax.set_zlim(mid[2] - r, mid[2] + r)
 
-        # Show
-        plt.show()
+        if should_return:
+            return self.ax
+        else:
+            # Show
+            plt.show()
+            return None