rustytorch/crates/models/rtx-vision-advanced/src/autonomous/tracking.rs

//! Multi-object tracking for autonomous vehicles

use crate::{VisionResult, VisionError};
use crate::detection::three_d::BoundingBox3D;
use crate::autonomous::{TrackedObject, ObjectClass, TrajectoryPoint, Timestamp};
use std::collections::HashMap;

/// Multi-object tracker
pub struct MultiObjectTracker {
    tracks: HashMap<usize, Track>,
    next_track_id: usize,
    max_age: usize,
    iou_threshold: f32,
}

impl MultiObjectTracker {
    pub fn new() -> VisionResult<Self> {
        Ok(Self {
            tracks: HashMap::new(),
            next_track_id: 0,
            max_age: 30,
            iou_threshold: 0.3,
        })
    }
    
    /// Update tracker with new detections
    pub fn update(&mut self, detections: &[BoundingBox3D], timestamp: Timestamp) -> VisionResult<Vec<TrackedObject>> {
        // Age existing tracks
        for track in self.tracks.values_mut() {
            track.age += 1;
            track.time_since_update = timestamp.duration_since(&track.last_update) as f32;
        }
        
        // Associate detections with existing tracks
        let associations = self.associate_detections(detections);
        
        // Update matched tracks
        for (track_id, detection) in associations {
            if let Some(track) = self.tracks.get_mut(&track_id) {
                track.update_with_detection(detection, timestamp);
            }
        }
        
        // Create new tracks for unmatched detections
        self.create_new_tracks(detections, timestamp);
        
        // Remove old tracks
        self.tracks.retain(|_, track| track.age <= self.max_age);
        
        // Convert tracks to tracked objects
        let tracked_objects = self.tracks.values()
            .map(|track| track.to_tracked_object())
            .collect();
        
        Ok(tracked_objects)
    }
    
    fn associate_detections(&self, detections: &[BoundingBox3D]) -> HashMap<usize, &BoundingBox3D> {
        let mut associations = HashMap::new();
        
        for detection in detections {
            let mut best_iou = 0.0;
            let mut best_track_id = None;
            
            for (&track_id, track) in &self.tracks {
                if let Some(ref last_bbox) = track.last_bbox {
                    let iou = detection.iou_3d(last_bbox);
                    if iou > best_iou && iou > self.iou_threshold {
                        best_iou = iou;
                        best_track_id = Some(track_id);
                    }
                }
            }
            
            if let Some(track_id) = best_track_id {
                associations.insert(track_id, detection);
            }
        }
        
        associations
    }
    
    fn create_new_tracks(&mut self, detections: &[BoundingBox3D], timestamp: Timestamp) {
        for detection in detections {
            let track = Track::new(self.next_track_id, detection.clone(), timestamp);
            self.tracks.insert(self.next_track_id, track);
            self.next_track_id += 1;
        }
    }
}

/// Individual track
struct Track {
    id: usize,
    last_bbox: Option<BoundingBox3D>,
    history: Vec<BoundingBox3D>,
    velocity: [f32; 3],
    age: usize,
    hits: usize,
    last_update: Timestamp,
}

impl Track {
    fn new(id: usize, initial_detection: BoundingBox3D, timestamp: Timestamp) -> Self {
        Self {
            id,
            last_bbox: Some(initial_detection.clone()),
            history: vec![initial_detection],
            velocity: [0.0, 0.0, 0.0],
            age: 0,
            hits: 1,
            last_update: timestamp,
        }
    }
    
    fn update_with_detection(&mut self, detection: &BoundingBox3D, timestamp: Timestamp) {
        // Update velocity estimate
        if let Some(ref last_bbox) = self.last_bbox {
            let dt = timestamp.duration_since(&self.last_update) as f32;
            if dt > 0.0 {
                for i in 0..3 {
                    self.velocity[i] = (detection.center[i] - last_bbox.center[i]) / dt;
                }
            }
        }
        
        self.last_bbox = Some(detection.clone());
        self.history.push(detection.clone());
        self.age = 0;
        self.hits += 1;
        self.last_update = timestamp;
        
        // Keep limited history
        if self.history.len() > 10 {
            self.history.remove(0);
        }
    }
    
    fn to_tracked_object(&self) -> TrackedObject {
        let bbox_3d = self.last_bbox.as_ref().unwrap().clone();
        let object_class = self.classify_object(&bbox_3d);
        
        TrackedObject {
            track_id: self.id,
            bbox_3d,
            velocity: self.velocity,
            track_confidence: self.calculate_track_confidence(),
            object_class,
            age: self.age,
            hits: self.hits,
            time_since_update: 0.0, // Will be updated by caller
            predicted_trajectory: self.predict_trajectory(),
        }
    }
    
    fn classify_object(&self, bbox_3d: &BoundingBox3D) -> ObjectClass {
        // Simple size-based classification
        let volume = bbox_3d.volume();
        
        if volume > 20.0 {
            ObjectClass::Vehicle
        } else if volume < 1.0 {
            ObjectClass::Pedestrian
        } else {
            ObjectClass::Unknown
        }
    }
    
    fn calculate_track_confidence(&self) -> f32 {
        // Confidence based on track stability
        let hit_ratio = self.hits as f32 / (self.age + 1) as f32;
        hit_ratio.min(1.0)
    }
    
    fn predict_trajectory(&self) -> Vec<TrajectoryPoint> {
        let mut trajectory = Vec::new();
        
        if let Some(ref current_bbox) = self.last_bbox {
            // Predict future positions using constant velocity model
            for t in 1..6 { // 5 future time steps
                let dt = t as f32 * 0.1; // 0.1 second intervals
                
                let position = [
                    current_bbox.center[0] + self.velocity[0] * dt,
                    current_bbox.center[1] + self.velocity[1] * dt,
                    current_bbox.center[2] + self.velocity[2] * dt,
                ];
                
                trajectory.push(TrajectoryPoint {
                    position,
                    time_offset: dt,
                    covariance: [[0.1, 0.0, 0.0], [0.0, 0.1, 0.0], [0.0, 0.0, 0.1]], // Simple diagonal covariance
                });
            }
        }
        
        trajectory
    }
}