follow camera condition

rostok/simulation_chrono/simulation.py

@@ -44,13 +44,14 @@ def chrono_NSC_system(solver_iterations=100, force_tolerance=1e-4, gravity_list=
 
 class ChronoVisManager():
 
-    def __init__(self, fps=100, delay: bool = False):
+    def __init__(self, fps=100, delay: bool = False, is_follow_camera = False):
         self.vis = chronoirr.ChVisualSystemIrrlicht()
         self.delay_flag = delay
         self.fps = fps
         self.step_counter = 0
+        self.is_follow_camera = is_follow_camera
 
-    def initialize_vis(self, chrono_system):
+    def initialize_vis(self, chrono_system, observed_body: chrono.ChBody = None):
         self.vis.AttachSystem(chrono_system)
         self.vis.SetWindowSize(1024, 768)
         self.vis.SetWindowTitle('Grab demo')
@@ -65,7 +66,10 @@ def initialize_vis(self, chrono_system):
         self.vis.AddLight(chrono.ChVectorD(0.3, 0.4, 0.3), 0.5, chrono.ChColor(0.7, 0.7, 0.7))
         self.vis.AddLight(chrono.ChVectorD(-0.3, 0.4, -0.3), 0.5, chrono.ChColor(0.7, 0.7, 0.7))
         self.vis.AddLight(chrono.ChVectorD(-0.3, 0.4, 0.3), 0.5, chrono.ChColor(0.7, 0.7, 0.7))
+        if observed_body:
+            self.bind_camera_to_body(observed_body)
         # self.vis.AddTypicalLights()
+
     def enable_collision_shape(self):
         self.vis.EnableCollisionShapeDrawing(True)
 
@@ -74,16 +78,35 @@ def visualization_step(self, step_length):
             self.step_counter = 0
             self.vis.Run()
             self.vis.BeginScene(True, True, chrono.ChColor(0.1, 0.1, 0.1))
+            if self.is_follow_camera:
+                self.calculate_camera_position()
             self.vis.Render()
             self.vis.EndScene()
+
             # just to slow down the simulation
             if self.delay_flag:
                 time.sleep(0.01)
 
         else:
             self.step_counter += 1
 
+    def bind_camera_to_body(self, ch_body: chrono.ChBody):
+        self.watched_body = ch_body
 
+    def calculate_camera_position(self):
+
+        # Setting camera in coordinate frame of Body
+        camera_cord = chrono.ChCoordsysD(chrono.ChVectorD(0.5, 0.5, 0))
+
+        transform = self.watched_body.GetFrame_COG_to_abs()
+        position = self.watched_body.GetPos()
+
+        camera_cord_global = camera_cord * transform.coord
+
+        self.watched_body.GetCoord()
+        self.vis.SetCameraTarget(position)
+        self.vis.SetCameraPosition(camera_cord_global.pos)
+
 class EnvCreator():
     """Setup environment for robot simulation.
 
@@ -283,8 +306,10 @@ def simulate(
                 flag_container: container of flags that controls simulation
                 visualize (bool): determine if run the visualization """
         self.initialize(number_of_steps)
+        # Select observed body for camera
+        observed_body  = self.chrono_system.Get_bodylist()[-1]
         if visualize:
-            self.vis_manager.initialize_vis(self.chrono_system)
+            self.vis_manager.initialize_vis(self.chrono_system, observed_body)
 
         self.vis_manager.fps = fps
         stop_flag = False