Update cone_detection.py

Author: Ren990420

GEMstack/onboard/perception/cone_detection.py

@@ -1,4 +1,4 @@
-from ...state import AllState, VehicleState, ObjectPose, ObjectFrameEnum, ObstacleState, ObstacleMaterialEnum, \
+from ...state import AllState, VehicleState, ObjectPose, ObjectFrameEnum, Obstacle, ObstacleMaterialEnum, \
     ObstacleStateEnum
 from ..interface.gem import GEMInterface
 from ..component import Component
@@ -146,7 +146,7 @@ def undistort_image(self, image, K, D):
         # print('--------undistort', end-start)
         return undistorted, newK
 
-    def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
+    def update(self, vehicle: VehicleState) -> Dict[str, Obstacle]:
         downsample = False
         # Gate guards against data not being present for both sensors:
         if self.latest_image is None or self.latest_lidar is None:
@@ -221,18 +221,18 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
         # Visualize the received images in 2D with their corresponding labels
         # It draws rectangles and labels on the images:
         if getattr(self, 'visualize_2d', False):
-            for (cx, cy, w, h, activity) in combined_boxes:
+            for (cx, cy, w, h, state) in combined_boxes:
                 left = int(cx - w / 2)
                 right = int(cx + w / 2)
                 top = int(cy - h / 2)
                 bottom = int(cy + h / 2)
-                if activity == ObstacleStateEnum.STANDING:
+                if state == ObstacleStateEnum.STANDING:
                     color = (255, 0, 0)
                     label = "STANDING"
-                elif activity == ObstacleStateEnum.RIGHT:
+                elif state == ObstacleStateEnum.RIGHT:
                     color = (0, 255, 0)
                     label = "RIGHT"
-                elif activity == ObstacleStateEnum.LEFT:
+                elif state == ObstacleStateEnum.LEFT:
                     color = (0, 0, 255)
                     label = "LEFT"
                 else:
@@ -259,7 +259,7 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
         obstacles = {}
 
         for i, box_info in enumerate(combined_boxes):
-            cx, cy, w, h, activity = box_info
+            cx, cy, w, h, state = box_info
             # print(cx, cy, w, h)
             left = int(cx - w / 1.6)
             right = int(cx + w / 1.6)
@@ -355,12 +355,12 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
                             roll=avg_roll,
                             frame=new_pose.frame
                         )
-                        updated_obstacle = ObstacleState(
+                        updated_obstacle = Obstacle(
                             pose=updated_pose,
                             dimensions=dims,
                             outline=None,
-                            type=ObstacleMaterialEnum.TRAFFIC_CONE,
-                            activity=activity,
+                            material=ObstacleMaterialEnum.TRAFFIC_CONE,
+                            state=state,
                         )
                     else:
                         updated_obstacle = old_state
@@ -369,24 +369,24 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
                 else:
                     obstacle_id = f"Cone{self.cone_counter}"
                     self.cone_counter += 1
-                    new_obstacle = ObstacleState(
+                    new_obstacle = Obstacle(
                         pose=new_pose,
                         dimensions=dims,
                         outline=None,
-                        type=ObstacleMaterialEnum.TRAFFIC_CONE,
-                        activity=activity,
+                        material=ObstacleMaterialEnum.TRAFFIC_CONE,
+                        state=state,
                     )
                     obstacles[obstacle_id] = new_obstacle
                     self.tracked_obstacles[obstacle_id] = new_obstacle
             else:
                 obstacle_id = f"Cone{self.cone_counter}"
                 self.cone_counter += 1
-                new_obstacle = ObstacleState(
+                new_obstacle = Obstacle(
                     pose=new_pose,
                     dimensions=dims,
                     outline=None,
-                    type=ObstacleMaterialEnum.TRAFFIC_CONE,
-                    activity=activity,
+                    material=ObstacleMaterialEnum.TRAFFIC_CONE,
+                    state=state,
                 )
                 obstacles[obstacle_id] = new_obstacle
 
@@ -400,7 +400,7 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
                     f"Cone ID: {obstacle_id}\n"
                     f"Pose: (x: {p.x:.3f}, y: {p.y:.3f}, z: {p.z:.3f}, "
                     f"yaw: {p.yaw:.3f}, pitch: {p.pitch:.3f}, roll: {p.roll:.3f})\n"
-                    f"type:{obstacle.activity}"
+                    f"state:{obstacle.state}"
                 )
             end = time.time()
             # print('-------processing time', end -start)
@@ -418,7 +418,7 @@ def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
                     f"Cone ID: {obstacle_id}\n"
                     f"Pose: (x: {p.x:.3f}, y: {p.y:.3f}, z: {p.z:.3f}, "
                     f"yaw: {p.yaw:.3f}, pitch: {p.pitch:.3f}, roll: {p.roll:.3f})\n"
-                    f"type:{obstacle.activity}"
+                    f"state:{obstacle.state}"
                 )
         end = time.time()
         # print('-------processing time', end -start)
@@ -482,7 +482,7 @@ def state_inputs(self):
     def state_outputs(self):
         return ['obstacles']
 
-    def update(self, vehicle: VehicleState) -> Dict[str, ObstacleState]:
+    def update(self, vehicle: VehicleState) -> Dict[str, Obstacle]:
         if self.t_start is None:
             self.t_start = self.vehicle_interface.time()
         t = self.vehicle_interface.time() - self.t_start
@@ -498,8 +498,8 @@ def box_to_fake_obstacle(box):
     x, y, w, h = box
     pose = ObjectPose(t=0, x=x + w / 2, y=y + h / 2, z=0, yaw=0, pitch=0, roll=0, frame=ObjectFrameEnum.CURRENT)
     dims = (w, h, 0)
-    return ObstacleState(pose=pose, dimensions=dims, outline=None,
-                         type=ObstacleMaterialEnum.TRAFFIC_CONE, activity=ObstacleStateEnum.STANDING)
+    return Obstacle(pose=pose, dimensions=dims, outline=None,
+                         material=ObstacleMaterialEnum.TRAFFIC_CONE, state=ObstacleStateEnum.STANDING)
 
 
 if __name__ == '__main__':