From 542a01de4d5c3480c17f21dbaf10dc1f2969f339 Mon Sep 17 00:00:00 2001
From: Michel Breyer <10465414+mbreyer@users.noreply.github.com>
Date: Fri, 3 Dec 2021 16:51:13 +0100
Subject: [PATCH] Visualize hw bbox

---
 scripts/hw_node.py | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/scripts/hw_node.py b/scripts/hw_node.py
index f8c161b..b9f3d77 100755
--- a/scripts/hw_node.py
+++ b/scripts/hw_node.py
@@ -5,6 +5,7 @@ import numpy as np
 import rospy
 
 from active_grasp.bbox import AABBox, to_bbox_msg
+from active_grasp.rviz import Visualizer
 from active_grasp.srv import *
 from robot_helpers.io import load_yaml
 from robot_helpers.ros.moveit import MoveItClient
@@ -16,13 +17,13 @@ class HwNode:
     def __init__(self):
         self.load_parameters()
         self.init_robot_connection()
+        self.init_visualizer()
         self.advertise_services()
         rospy.spin()
 
     def load_parameters(self):
-        cfg = rospy.get_param("hw")
-        self.T_base_roi = Transform.from_matrix(np.loadtxt(cfg["roi_calib_file"]))
-        self.scene_config = load_yaml(cfg["scene_file"])
+        self.cfg = rospy.get_param("hw")
+        self.T_base_roi = Transform.from_matrix(np.loadtxt(self.cfg["roi_calib_file"]))
 
     def init_robot_connection(self):
         self.gripper = PandaGripperClient()
@@ -31,6 +32,10 @@ class HwNode:
         )
         self.moveit = MoveItClient("panda_arm")
 
+    def init_visualizer(self):
+        self.vis = Visualizer()
+        rospy.Timer(rospy.Duration(1), self.draw_bbox)
+
     def advertise_services(self):
         rospy.Service("seed", Seed, self.seed)
         rospy.Service("reset", Reset, self.reset)
@@ -40,18 +45,23 @@ class HwNode:
         return SeedResponse()
 
     def reset(self, req):
+        q0, bbox = self.load_config()
+
         # Move to the initial configuration
         self.switch_to_joint_trajectory_controller()
-        self.moveit.goto(self.scene_config["q0"])
+        self.moveit.goto(q0)
         self.gripper.move(0.04)
 
-        # Construct bounding box
-        bbox_min = self.T_base_roi.apply(self.scene_config["target"]["min"])
-        bbox_max = self.T_base_roi.apply(self.scene_config["target"]["max"])
-        bbox = AABBox(bbox_min, bbox_max)
-
         return ResetResponse(to_bbox_msg(bbox))
 
+    def load_config(self):
+        scene_config = load_yaml(self.cfg["scene_file"])
+        q0 = scene_config["q0"]
+        bbox_min = self.T_base_roi.apply(scene_config["target"]["min"])
+        bbox_max = self.T_base_roi.apply(scene_config["target"]["max"])
+        bbox = AABBox(bbox_min, bbox_max)
+        return q0, bbox
+
     def switch_to_joint_trajectory_controller(self):
         req = SwitchControllerRequest()
         req.start_controllers = ["position_joint_trajectory_controller"]
@@ -59,6 +69,10 @@ class HwNode:
         req.strictness = 1
         self.switch_controller(req)
 
+    def draw_bbox(self, event):
+        _, bbox = self.load_config()
+        self.vis.bbox("panda_link0", bbox)
+
 
 def main():
     rospy.init_node("hw")