Added some comments to explain complicated sections of code and moved data saving functionality to seperate function

Author: LucasEby

GEMstack/onboard/perception/cone_detection.py

@@ -157,61 +157,29 @@ def undistort_image(self, image, K, D):
         return undistorted, newK
 
     def update(self, vehicle: VehicleState) -> Dict[str, AgentState]:
-
-        start = time.time()
         downsample = False
+        # Gate guards against data not being present for both sensors:
         if self.latest_image is None or self.latest_lidar is None:
             return {}
         lastest_image = self.latest_image.copy()
-        # Ensure data/ exists and build timestamp
+        
+        # Set up current time variables
+        start = time.time()
+        current_time = self.vehicle_interface.time()
+
         if downsample:
             lidar_down = downsample_points(self.latest_lidar, voxel_size=0.1)
         else:
             lidar_down = self.latest_lidar.copy()
-        current_time = self.vehicle_interface.time()
+
         if self.start_time is None:
             self.start_time = current_time
         time_elapsed = current_time - self.start_time
 
+        # Ensure data/ exists and build timestamp
         if self.save_data:
-            os.makedirs("data", exist_ok=True)
-            tstamp = int(self.vehicle_interface.time() * 1000)
-            # 1) Dump raw image
-            cv2.imwrite(f"data/{tstamp}_image.png", lastest_image)
-            # 2) Dump raw LiDAR
-            np.savez(f"data/{tstamp}_lidar.npz", lidar=self.latest_lidar)
-            # 3) Write BEFORE_TRANSFORM
-            with open(f"data/{tstamp}_vehstate.txt", "w") as f:
-                vp = vehicle.pose
-                f.write(
-                    f"BEFORE_TRANSFORM "
-                    f"x={vp.x:.3f}, y={vp.y:.3f}, z={vp.z:.3f}, "
-                    f"yaw={vp.yaw:.2f}, pitch={vp.pitch:.2f}, roll={vp.roll:.2f}\n"
-                )
-            # Compute vehicle_start_pose in either START or CURRENT
-            if self.use_start_frame:
-                if self.start_pose_abs is None:
-                    self.start_pose_abs = vehicle.pose
-                vehicle_start_pose = vehicle.pose.to_frame(
-                    ObjectFrameEnum.START,
-                    vehicle.pose,
-                    self.start_pose_abs
-                )
-                mode = "START"
-            else:
-                vehicle_start_pose = vehicle.pose
-                mode = "CURRENT"
-            with open(f"data/{tstamp}_vehstate.txt", "a") as f:
-                f.write(
-                    f"AFTER_TRANSFORM "
-                    f"x={vehicle_start_pose.x:.3f}, "
-                    f"y={vehicle_start_pose.y:.3f}, "
-                    f"z={vehicle_start_pose.z:.3f}, "
-                    f"yaw={vehicle_start_pose.yaw:.2f}, "
-                    f"pitch={vehicle_start_pose.pitch:.2f}, "
-                    f"roll={vehicle_start_pose.roll:.2f}, "
-                    f"frame={mode}\n"
-                )
+            self.save_sensor_data(vehicle=vehicle, latest_image=latest_image)
+        
         if self.camera_front == False:
             start = time.time()
             undistorted_img, current_K = self.undistort_image(lastest_image, self.K, self.D)
@@ -254,6 +222,8 @@ def update(self, vehicle: VehicleState) -> Dict[str, AgentState]:
             cx, cy, w, h = box
             combined_boxes.append((cx, cy, w, h, AgentActivityEnum.STANDING))
 
+        # 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:
                 left = int(cx - w / 2)
@@ -278,8 +248,15 @@ def update(self, vehicle: VehicleState) -> Dict[str, AgentState]:
             cv2.imshow("Detection - Cone 2D", undistorted_img)
 
         start = time.time()
+        # Transform the lidar points from lidar frame of reference to camera EXTRINSIC frame of reference.
+        # Then project the pixels onto the lidar points to "paint them" (essentially determine which points are associated with detected objects)
         pts_cam = transform_points_l2c(lidar_down, self.T_l2c)
         projected_pts = project_points(pts_cam, self.current_K, lidar_down)
+        # What is returned:
+        # projected_pts[:, 0]: u-coordinate in the image (horizontal pixel position)
+        # projected_pts[:, 1]: v-coordinate in the image (vertical pixel position)
+        # projected_pts[:, 2:5]: original X, Y, Z coordinates in the LiDAR frame
+
         end = time.time()
         # print('-------processing time1---', end -start)
 
@@ -456,6 +433,46 @@ def update(self, vehicle: VehicleState) -> Dict[str, AgentState]:
         # print('-------processing time', end -start)
         return self.tracked_agents
 
+    def save_sensor_data(self, vehicle: VehicleState, latest_image) -> None:
+        os.makedirs("data", exist_ok=True)
+        tstamp = int(self.vehicle_interface.time() * 1000)
+        # 1) Dump raw image
+        cv2.imwrite(f"data/{tstamp}_image.png", lastest_image)
+        # 2) Dump raw LiDAR
+        np.savez(f"data/{tstamp}_lidar.npz", lidar=self.latest_lidar)
+        # 3) Write BEFORE_TRANSFORM
+        with open(f"data/{tstamp}_vehstate.txt", "w") as f:
+            vp = vehicle.pose
+            f.write(
+                f"BEFORE_TRANSFORM "
+                f"x={vp.x:.3f}, y={vp.y:.3f}, z={vp.z:.3f}, "
+                f"yaw={vp.yaw:.2f}, pitch={vp.pitch:.2f}, roll={vp.roll:.2f}\n"
+            )
+        # Compute vehicle_start_pose in either START or CURRENT
+        if self.use_start_frame:
+            if self.start_pose_abs is None:
+                self.start_pose_abs = vehicle.pose
+            vehicle_start_pose = vehicle.pose.to_frame(
+                ObjectFrameEnum.START,
+                vehicle.pose,
+                self.start_pose_abs
+            )
+            mode = "START"
+        else:
+            vehicle_start_pose = vehicle.pose
+            mode = "CURRENT"
+        with open(f"data/{tstamp}_vehstate.txt", "a") as f:
+            f.write(
+                f"AFTER_TRANSFORM "
+                f"x={vehicle_start_pose.x:.3f}, "
+                f"y={vehicle_start_pose.y:.3f}, "
+                f"z={vehicle_start_pose.z:.3f}, "
+                f"yaw={vehicle_start_pose.yaw:.2f}, "
+                f"pitch={vehicle_start_pose.pitch:.2f}, "
+                f"roll={vehicle_start_pose.roll:.2f}, "
+                f"frame={mode}\n"
+            )
+
     # ----- Fake Cone Detector 2D (for Testing Purposes) -----