Custom VAE for trajectroy prediction

trap/base.py

@@ -308,6 +308,37 @@ class Track:
             pos += distance
         
         return new_path
+    
+    def get_simplified_history_with_absolute_distance(self, distance: float, camera: Camera) -> list[tuple[float, float]]:
+        # Similar to get_simplified_history, but with absolute world-space distance
+        # not the distance of the track length
+
+        if len(self.history) < 1:
+            return []
+
+    
+        path = self.get_projected_history(H=None, camera=camera)
+        new_path: List[dict] = [path[0]]
+
+        distance_sq = distance**2
+        
+        for a, b in zip(path[:-1], path[1:]):
+            # check if segment has our next point (pos)
+            # because running sequentially, this is if  point b
+            # is lower then our target position
+            b_distance_sq = ((b[0]-new_path[0])**2 + (b[1]-new_path[1])**2)
+
+            if b_distance_sq <= distance_sq: 
+                continue
+            
+            a_distance_sq = ((a[0]-new_path[0])**2 + (a[1]-new_path[1])**2)
+
+            relative_t = inv_lerp(a_distance_sq, b_distance_sq, distance_sq)
+            x = lerp(a[0], b[0], relative_t)
+            y = lerp(a[1], b[1], relative_t)
+            new_path.append([x,y])
+        
+        return new_path
                 
 
 
@@ -406,6 +437,15 @@ class Track:
             }
 
         return pd.DataFrame(data_dict, columns=data_columns)
+    
+    def to_flat_dataframe(self, camera: Camera) -> pd.DataFrame:
+        positions = self.get_projected_history(None, camera)
+        data = pd.DataFrame(positions, columns=['x', 'y'])
+
+        data['dx'] = data['x'].diff()
+        data['dy'] = data['y'].diff()
+        
+        return data.bfill()
 
     def to_trajectron_node(self, camera: Camera, env: Environment) -> Node:
         node_data = self.to_dataframe(camera)