Pass frame explicitly to each visualization method

This commit is contained in:
Michel Breyer
2021-08-17 16:24:23 +02:00
parent 69e6acf00b
commit 26f7103237
4 changed files with 53 additions and 36 deletions

View File

@@ -36,7 +36,7 @@ class BasePolicy(Policy):
self.score_fn = lambda g: g.pose.translation[2] # TODO
def init_visualizer(self):
self.visualizer = Visualizer(self.base_frame)
self.visualizer = Visualizer()
def activate(self, bbox):
self.bbox = bbox
@@ -61,11 +61,12 @@ class BasePolicy(Policy):
self.best_grasp = None
self.visualizer.clear()
self.visualizer.bbox(bbox)
self.visualizer.bbox(self.base_frame, bbox)
def integrate_img(self, img, extrinsic):
self.viewpoints.append(extrinsic.inv())
self.tsdf.integrate(img, self.intrinsic, extrinsic * self.T_base_task)
tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
if self.filter_grasps:
out = self.vgn.predict(self.tsdf.get_grid())
@@ -75,10 +76,11 @@ class BasePolicy(Policy):
self.width_hist[t, ...] = out.width
mean_qual = self.compute_mean_quality()
self.visualizer.quality(self.task_frame, self.tsdf.voxel_size, mean_qual)
self.visualizer.quality(self.task_frame, voxel_size, mean_qual)
self.visualizer.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
self.visualizer.path(self.viewpoints)
self.visualizer.map_cloud(self.task_frame, voxel_size, tsdf_grid)
self.visualizer.path(self.base_frame, self.viewpoints)
def compute_best_grasp(self):
if self.filter_grasps:
@@ -96,7 +98,7 @@ class BasePolicy(Policy):
grasps = sort_grasps(grasps, self.score_fn)
self.visualizer.quality(self.task_frame, self.tsdf.voxel_size, qual)
self.visualizer.grasps(grasps)
self.visualizer.grasps(self.base_frame, grasps)
return grasps[0] if len(grasps) > 0 else None