From fcee4a9517d7e41d12c5f9a9e37e95e9c78b9d10 Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Fri, 8 Oct 2021 17:19:36 +0200
Subject: [PATCH] Reuse box lines computation

---
 src/active_grasp/policy.py     |  1 +
 src/active_grasp/rviz.py       | 28 ++++++----------------------
 src/active_grasp/simulation.py |  2 +-
 3 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/src/active_grasp/policy.py b/src/active_grasp/policy.py
index adf1af1..3e3856f 100644
--- a/src/active_grasp/policy.py
+++ b/src/active_grasp/policy.py
@@ -63,6 +63,7 @@ class Policy:
         self.T_task_base = self.T_base_task.inv()
         tf.broadcast(self.T_base_task, self.base_frame, self.task_frame)
         rospy.sleep(1.0)  # Wait for tf tree to be updated
+        self.vis.workspace(self.task_frame, 0.3)
 
     def update(self, img, x, q):
         raise NotImplementedError
diff --git a/src/active_grasp/rviz.py b/src/active_grasp/rviz.py
index 7cbae1e..1a1e8eb 100644
--- a/src/active_grasp/rviz.py
+++ b/src/active_grasp/rviz.py
@@ -3,10 +3,11 @@ import numpy as np
 from robot_helpers.ros.rviz import *
 from robot_helpers.spatial import Transform
 import vgn.rviz
+from vgn.utils import box_lines
 
 cm = lambda s: tuple([float(1 - s), float(s), float(0)])
-red = np.r_[1.0, 0.0, 0.0]
-blue = np.r_[0, 0.6, 1.0]
+red = [1.0, 0.0, 0.0]
+blue = [0, 0.6, 1.0]
 
 
 class Visualizer(vgn.rviz.Visualizer):
@@ -14,22 +15,7 @@ class Visualizer(vgn.rviz.Visualizer):
         pose = Transform.identity()
         scale = [0.004, 0.0, 0.0]
         color = red
-        corners = bbox.corners
-        edges = [
-            (0, 1),
-            (1, 3),
-            (3, 2),
-            (2, 0),
-            (4, 5),
-            (5, 7),
-            (7, 6),
-            (6, 4),
-            (0, 4),
-            (1, 5),
-            (3, 7),
-            (2, 6),
-        ]
-        lines = [(corners[s], corners[e]) for s, e in edges]
+        lines = box_lines(bbox.min, bbox.max)
         marker = create_line_list_marker(frame, pose, scale, color, lines, "bbox")
         self.draw([marker])
 
@@ -88,7 +74,7 @@ class Visualizer(vgn.rviz.Visualizer):
         markers = []
         for i, (view, value) in enumerate(zip(views, values)):
             color = cm((value - vmin) / (vmax - vmin))
-            marker = create_cam_view_marker(
+            marker = create_view_marker(
                 frame,
                 view,
                 scale,
@@ -103,9 +89,7 @@ class Visualizer(vgn.rviz.Visualizer):
         self.draw(markers)
 
 
-def create_cam_view_marker(
-    frame, pose, scale, color, intrinsic, near, far, ns="", id=0
-):
+def create_view_marker(frame, pose, scale, color, intrinsic, near, far, ns="", id=0):
     marker = create_marker(Marker.LINE_LIST, frame, pose, scale, color, ns, id)
     x_n = near * intrinsic.width / (2.0 * intrinsic.fx)
     y_n = near * intrinsic.height / (2.0 * intrinsic.fy)
diff --git a/src/active_grasp/simulation.py b/src/active_grasp/simulation.py
index 9ac288e..2a95923 100644
--- a/src/active_grasp/simulation.py
+++ b/src/active_grasp/simulation.py
@@ -170,7 +170,7 @@ class CustomScene(Scene):
 
     def load_config(self):
         with self.config_path.open("r") as f:
-            self.scene = yaml.load(f)
+            self.scene = yaml.load(f, Loader=yaml.FullLoader)
         self.center = np.asarray(self.scene["center"])
         self.length = 0.3
         self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0]