From 7a53fddc314a5d1310d9cbc74e5e2a819238e611 Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Thu, 26 Aug 2021 11:43:03 +0200
Subject: [PATCH] Minor

---
 active_grasp/nbv.py | 64 ++++++++++++++++++++++-----------------------
 1 file changed, 31 insertions(+), 33 deletions(-)

diff --git a/active_grasp/nbv.py b/active_grasp/nbv.py
index f9b0d24..257e185 100644
--- a/active_grasp/nbv.py
+++ b/active_grasp/nbv.py
@@ -1,47 +1,34 @@
 import itertools
 import numpy as np
-import rospy
 
-from .policy import BasePolicy
+from .policy import MultiViewPolicy
 from vgn.utils import look_at, spherical_to_cartesian
 
 
-class NextBestView(BasePolicy):
-    def __init__(self, rate, filter_grasps):
-        super().__init__(rate, filter_grasps)
-        self.max_viewpoints = 20
-        self.min_gain = 10.0
-
-    def activate(self, bbox):
-        super().activate(bbox)
+class NextBestView(MultiViewPolicy):
+    def __init__(self, rate):
+        super().__init__(rate)
+        self.max_views = 20
+        self.min_ig = 10.0
 
     def update(self, img, extrinsic):
-        # Integrate latest measurement
-        self.integrate_img(img, extrinsic)
-
-        # Generate viewpoints
-        views = self.generate_viewpoints()
-
-        # Evaluate viewpoints
-        gains = [self.compute_ig(v) for v in views]
-        costs = [self.compute_cost(v) for v in views]
+        self.integrate(img, extrinsic)
+        views = self.generate_views()
+        gains = self.compute_expected_information_gains(views)
+        costs = self.compute_movement_costs(views)
         utilities = gains / np.sum(gains) - costs / np.sum(costs)
 
-        # Visualize
         self.vis.views(self.base_frame, self.intrinsic, views, utilities)
 
-        # Determine next-best-view
         i = np.argmax(utilities)
-        nbv, g = views[i], gains[i]
+        nbv, ig = views[i], gains[i]
 
-        if self.check_done(g):
-            self.best_grasp = self.compute_best_grasp()
+        if self.check_stopping_criteria(ig):
             self.done = True
-            return
+        else:
+            return nbv.inv()  # the controller expects T_cam_base
 
-        return nbv.inv()  # the controller expects T_cam_base
-
-    def generate_viewpoints(self):
+    def generate_views(self):
         r, h = 0.14, 0.2
         thetas = np.arange(1, 4) * np.deg2rad(30)
         phis = np.arange(1, 6) * np.deg2rad(60)
@@ -53,7 +40,21 @@ class NextBestView(BasePolicy):
             views.append(look_at(eye, target, up).inv())
         return views
 
-    def compute_ig(self, view, downsample=20):
+    def compute_expected_information_gains(self, views):
+        return [self.ig_fn(v) for v in views]
+
+    def compute_movement_costs(self, views):
+        return [self.cost_fn(v) for v in views]
+
+    def check_stopping_criteria(self, ig):
+        if len(self.views) > self.max_views:
+            return True
+        if ig < self.min_ig:
+            return True
+        # TODO check whether a "good enough" grasp has been found
+        return False
+
+    def ig_fn(self, view, downsample=20):
         fx = self.intrinsic.fx / downsample
         fy = self.intrinsic.fy / downsample
         cx = self.intrinsic.cx / downsample
@@ -114,8 +115,5 @@ class NextBestView(BasePolicy):
 
         return ig
 
-    def compute_cost(self, view):
+    def cost_fn(self, view):
         return 1.0
-
-    def check_done(self, gain):
-        return len(self.viewpoints) > self.max_viewpoints or gain < self.min_gain