Putting this here for now

LucasEby · LucasEby · commit b51e87e022ae · 2025-02-18T01:46:53.000Z
diff --git a/GEMstack/onboard/perception/AgentTracker.py b/GEMstack/onboard/perception/AgentTracker.py
@@ -0,0 +1,144 @@
+import math
+from typing import Dict, List
+
+from GEMstack.onboard.perception.IdTracker import IdTracker
+from GEMstack.onboard.perception.PrevAgent import PrevAgent
+from GEMstack.state.agent import AgentState
+
+
+class AgentTracker():
+    """Associates and tracks AgentState agents.
+    """
+    def __init__(self):
+        # List of PrevAgent objects (each keeps track of the last seen state and time since seen)
+        self.prev_agents: List[PrevAgent] = []
+        # List of currently visible AgentState objects
+        self.current_agents: List[AgentState] = []
+        # Maximum time (in seconds) to keep a lost agent before dropping it.
+        self.drop_agent_t: float = 1.0
+        # Id tracker for creating new unique pedestrian IDs.
+        self.id_tracker = IdTracker()
+    
+    def assign_ids(self, agents: list) -> Dict[str,AgentState]:
+        # Act with the assumption that you are being sent a list of AgentState objects and you need to use the object fields to keep track of them for your task
+        # Further act on the assumption that we will decide the id's of the pedestrians by assuming that 2 pedestrians are the same pedestrian if a
+        # previously stored AgentState pose and dimensions overlap with a newly passed in AgentState
+        # Act on the assumption that the AgentState objects are all in reference to the start frame of the vehicle
+        # some helper functions in this class, LostAgent.py, and IdTracker.py have been created to try to help you out with your task.
+        # Assume that the output returned from this function will be a dictionary of AgentState objects with the key corresponding to their id
+        """
+        Associates new AgentState objects with existing tracked agents based on overlap.
+        If an agent does not match any previously tracked agent, a new unique id is assigned.
+        Also updates the “lost” time for agents that are not matched in this frame.
+
+        Parameters:
+            agents (list): List of AgentState objects for the current frame 
+                           (already converted to the start frame of reference).
+
+        Returns:
+            Dict[str, AgentState]: Dictionary mapping agent id (as a string) to AgentState.
+        """
+        dt = 1.0  # Assuming a fixed time-step of 1 second (for simplicity)
+        output_agents: Dict[str, AgentState] = {}
+        matched_ids = set()
+        updated_prev_agents: List[PrevAgent] = []
+
+        # Process each new agent from the current frame.
+        for new_agent in agents:
+            found_match = None
+            # Look for a previously tracked agent whose bounding box overlaps.
+            for prev in self.prev_agents:
+                if prev.last_id in matched_ids:
+                    continue  # Already matched with another new agent.
+                if self.__agents_overlap(prev.last_state, new_agent):
+                    found_match = prev
+                    break
+
+            if found_match is not None:
+                # update velocity using the change in position.
+                if hasattr(new_agent, 'velocity'):
+                    new_agent.velocity = self.__calculate_velocity(found_match.last_state, new_agent, dt)
+
+                # Update the matched agent’s state and reset its lost-time counter.
+                found_match.last_state = new_agent
+                found_match.time_since_seen = 0.0
+                matched_ids.add(found_match.last_id)
+                output_agents[str(found_match.last_id)] = new_agent
+                updated_prev_agents.append(found_match)
+            else:
+                # No match found – assign a new unique id.
+                new_id = self.id_tracker.get_new_ped_id()
+                new_prev = PrevAgent(new_id, new_agent)
+                # initialize velocity to 0.
+                if hasattr(new_agent, 'velocity'):
+                    new_agent.velocity = 0.0
+                output_agents[str(new_id)] = new_agent
+                updated_prev_agents.append(new_prev)
+
+        # For all previously tracked agents that were not matched this frame,
+        # update their time-since-seen and only keep them if they have not timed out.
+        for prev in self.prev_agents:
+            if prev.last_id not in matched_ids:
+                prev.update_time(dt)
+                if prev.time_since_seen < self.drop_agent_t:
+                    updated_prev_agents.append(prev)
+
+        # Save the updated list of tracked agents.
+        self.prev_agents = updated_prev_agents
+        # update current_agents to include only those seen in the current frame.
+        self.current_agents = list(output_agents.values())
+        return output_agents
+
+    def __convert_to_start_frame(self):
+        """Converts a list of AgentState agents from ouster Lidar frame of 
+        reference (which is in reference to the current frame) to start 
+        frame frame of reference
+        """
+        # you can ignore this function akul
+        pass
+        
+    def __agents_overlap(self, ped1, ped2) -> bool:
+        """
+        Determines if two AgentState objects overlap based on their pose and dimensions.
+
+        Assumes each AgentState has:
+            - pose with attributes x and y.
+            - dimensions with attributes width and height.
+
+        Returns:
+            bool: True if they overlap; False otherwise.
+        """
+        # Get first agent's properties.
+        x1, y1 = ped1.pose.x, ped1.pose.y
+        w1, h1 = ped1.dimensions.width, ped1.dimensions.height
+
+        # Get second agent's properties.
+        x2, y2 = ped2.pose.x, ped2.pose.y
+        w2, h2 = ped2.dimensions.width, ped2.dimensions.height
+
+        # Compute bounding boxes (assuming (x, y) is the center).
+        left1, right1 = x1 - w1 / 2, x1 + w1 / 2
+        top1, bottom1 = y1 - h1 / 2, y1 + h1 / 2
+
+        left2, right2 = x2 - w2 / 2, x2 + w2 / 2
+        top2, bottom2 = y2 - h2 / 2, y2 + h2 / 2
+
+        # Check for overlap between the bounding boxes.
+        overlap = not (right1 < left2 or left1 > right2 or bottom1 < top2 or top1 > bottom2)
+        return overlap
+
+    def __calculate_velocity(self, old_state: AgentState, new_state: AgentState, dt: float) -> float:
+        """
+        Calculates the velocity based on the change in pose over time.
+
+        Parameters:
+            old_state (AgentState): The previous state.
+            new_state (AgentState): The current state.
+            dt (float): Time difference between the two states.
+
+        Returns:
+            float: The computed velocity.
+        """
+        dx = new_state.pose.x - old_state.pose.x
+        dy = new_state.pose.y - old_state.pose.y
+        return math.sqrt(dx * dx + dy * dy) / dt if dt > 0 else 0.0
diff --git a/GEMstack/onboard/perception/IdTracker.py b/GEMstack/onboard/perception/IdTracker.py
@@ -0,0 +1,12 @@
+class IdTracker():
+    """Abstracts out id tracking
+    """
+    def __init__(self):
+        self.__ped_id = 0
+
+    def get_new_ped_id(self) -> int:
+        """Returns a unique pedestrian id
+        """
+        assigned_id = self.__id
+        self.__ped_id += 1 # id will intentionally overflow to get back to 0
+        return assigned_id
diff --git a/GEMstack/onboard/perception/PrevAgent.py b/GEMstack/onboard/perception/PrevAgent.py
@@ -0,0 +1,17 @@
+# GEM imports:
+from ...state import AgentState
+
+class PrevAgent():
+    def __init__(self, last_id: int, last_state: AgentState):
+        self.time_since_seen: float = 0.0 # Time since the agent was last seen in seconds
+        self.last_id: int = last_id
+        self.last_state: AgentState = last_state
+
+    def update_time(time: float):
+        """Updates the time since the agent was last seen
+        """
+        if time <= 0:
+            # TODO: log error here
+            print("UPDATE TIME FOR LOST AGENT WAS LESS THAN OR EQUAL TO 0")
+        else:
+            self.time_since_seen += time
diff --git a/GEMstack/onboard/perception/pedestrian_detection.py b/GEMstack/onboard/perception/pedestrian_detection.py
@@ -48,6 +48,8 @@
 from .pedestrian_detection_utils import *
 from ..interface.gem import GEMInterface
 from ..component import Component
+from AgentTracker import AgentTracker
+
 
 def box_to_fake_agent(box):
     """Creates a fake agent state from an (x,y,w,h) bounding box.
@@ -101,6 +103,8 @@ def __init__(self, vehicle_interface : GEMInterface) -> None:
         # TF listener to get transformation from LiDAR to Camera
         self.tf_listener = tf.TransformListener()
 
+        self.agent_tracker = AgentTracker()
+
         if self.debug: self.init_debug()
     
     def init_debug(self,) -> None:
@@ -136,21 +140,21 @@ def find_dims(self, clusters: List[List[np.ndarray]]) -> List[np.ndarray]:
     #            Work towards own tracking class instead of simple YOLO track?
     #            Fix division by time
     # ret: Dict[track_id: vel[x, y, z]]
-    def find_vels(self, track_ids: List[int], obj_centers: List[np.ndarray]) -> Dict[int, np.ndarray]:
-        # Object not seen -> velocity = None
-        track_id_center_map = dict(zip(track_ids, obj_centers))
-        vels = defaultdict(lambda: np.array(())) # None is faster, np.array matches other find_ methods.
-
-        for prev_track_id, prev_agent in self.prev_agents.items():
-            if prev_track_id in track_ids:
-                # TODO: Add prev_agents to memory to avoid None velocity
-                # We should only be missing prev pose on first sight of track_id Agent.
-                # print("shape 1: ", track_id_center_map[prev_agent.track_id])
-                # print("shape 2: ", np.array([prev_agent.pose.x, prev_agent.pose.y, prev_agent.pose.z]))
-                # prev can be 3 separate Nones, current is just empty array... make this symmetrical
-                if prev_agent.pose.x and prev_agent.pose.y and prev_agent.pose.z and track_id_center_map[prev_agent.track_id].shape == 3:
-                    vels[prev_track_id] = (track_id_center_map[prev_track_id] - np.array([prev_agent.pose.x, prev_agent.pose.y, prev_agent.pose.z])) / (self.curr_time - self.prev_time)
-        return vels
+    # def find_vels(self, track_ids: List[int], obj_centers: List[np.ndarray]) -> Dict[int, np.ndarray]:
+    #     # Object not seen -> velocity = None
+    #     track_id_center_map = dict(zip(track_ids, obj_centers))
+    #     vels = defaultdict(lambda: np.array(())) # None is faster, np.array matches other find_ methods.
+
+    #     for prev_track_id, prev_agent in self.prev_agents.items():
+    #         if prev_track_id in track_ids:
+    #             # TODO: Add prev_agents to memory to avoid None velocity
+    #             # We should only be missing prev pose on first sight of track_id Agent.
+    #             # print("shape 1: ", track_id_center_map[prev_agent.track_id])
+    #             # print("shape 2: ", np.array([prev_agent.pose.x, prev_agent.pose.y, prev_agent.pose.z]))
+    #             # prev can be 3 separate Nones, current is just empty array... make this symmetrical
+    #             if prev_agent.pose.x and prev_agent.pose.y and prev_agent.pose.z and track_id_center_map[prev_agent.track_id].shape == 3:
+    #                 vels[prev_track_id] = (track_id_center_map[prev_track_id] - np.array([prev_agent.pose.x, prev_agent.pose.y, prev_agent.pose.z])) / (self.curr_time - self.prev_time)
+    #     return vels
 
 
     # TODO: Separate debug/viz class, bbox and 2d 3d points funcs 
@@ -231,25 +235,37 @@ def update_object_states(self, track_result: List[Results], extracted_pts_all: L
         #       Separate numpy prob still faster for now
         obj_centers = self.find_centers(pedestrians_3d_pts)
         obj_dims = self.find_dims(pedestrians_3d_pts)
-        obj_vels = self.find_vels(track_ids, obj_centers)
+
+        print("obj centers shape:")
+        print(obj_centers.shape)
+
+        print("obj dims shape: ")
+        print(obj_dims.shape)
+
+        # Assign id's based on whether or not agents overlap:
+        # self.current_agents = self.agent_tracker.assign_ids(self.prev_agents, obj_centers, obj_dims)
+
+        
+
+        # obj_vels = self.find_vels(track_ids, obj_centers)
 
         # Update Current AgentStates
-        for ind in range(num_objs):
-            obj_center = (None, None, None) if obj_centers[ind].size == 0 else obj_centers[ind]
-            obj_dim = (None, None, None) if obj_dims[ind].size == 0 else obj_dims[ind]
-            self.current_agents[track_ids[ind]] = (
-                AgentState(
-                    track_id = track_ids[ind],
-                    pose=ObjectPose(t=0, x=obj_center[0], y=obj_center[1], z=obj_center[2] ,yaw=0,pitch=0,roll=0,frame=ObjectFrameEnum.CURRENT),
-                    # (l, w, h)
-                    # TODO: confirm (z -> l, x -> w, y -> h)
-                    dimensions=(obj_dim[0], obj_dim[1], obj_dim[2]),  
-                    outline=None,
-                    type=AgentEnum.PEDESTRIAN,
-                    activity=AgentActivityEnum.MOVING,
-                    velocity= None if obj_vels[track_ids[ind]].size == 0 else tuple(obj_vels[track_ids[ind]]),
-                    yaw_rate=0
-                ))
+        # for ind in range(num_objs):
+        #     obj_center = (None, None, None) if obj_centers[ind].size == 0 else obj_centers[ind]
+        #     obj_dim = (None, None, None) if obj_dims[ind].size == 0 else obj_dims[ind]
+        #     self.current_agents[track_ids[ind]] = (
+        #         AgentState(
+        #             track_id = track_ids[ind],
+        #             pose=ObjectPose(t=0, x=obj_center[0], y=obj_center[1], z=obj_center[2] ,yaw=0,pitch=0,roll=0,frame=ObjectFrameEnum.CURRENT),
+        #             # (l, w, h)
+        #             # TODO: confirm (z -> l, x -> w, y -> h)
+        #             dimensions=(obj_dim[0], obj_dim[1], obj_dim[2]),  
+        #             outline=None,
+        #             type=AgentEnum.PEDESTRIAN,
+        #             activity=AgentActivityEnum.MOVING,
+        #             velocity= None if obj_vels[track_ids[ind]].size == 0 else tuple(obj_vels[track_ids[ind]]),
+        #             yaw_rate=0
+        #         ))
 
     def ouster_oak_callback(self, rgb_image_msg: Image, lidar_pc2_msg: PointCloud2):
         # Convert to cv2 image and run detector
