""" Object tracking script that reads predictions.json and outputs tracking.json with track IDs. Usage: python track_objects.py --input runs/val_viz/exp/predictions.json --output runs/val_viz/exp/tracking.json """ import argparse import csv import json import os import re from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path import numpy as np try: import yaml except ImportError: yaml = None try: from scipy.optimize import linear_sum_assignment SCIPY_AVAILABLE = True except ImportError: SCIPY_AVAILABLE = False print("Warning: scipy not available, using greedy matching instead of Hungarian algorithm") from merge_tracking_results import merge_case, normalize_image_name DEFAULT_DATA_CONFIG_PATH = Path(__file__).resolve().parents[2] / 'ultralytics' / 'cfg' / 'datasets' / 'mono3d_ground.yaml' DEFAULT_CLASS_MAP = { 'car': 0, 'suv': 1, 'pickup': 2, 'medium_car': 3, 'van': 4, 'bus': 5, 'truck': 6, 'tanker': 6, 'large_truck': 6, 'construction_vehicle': 6, 'special_vehicle': 7, 'unknown': 8, 'pedestrian': 9, 'bicyclist': 10, 'motorcyclist': 10, 'bicycle': 11, 'motorcycle': 11, 'tricycle': 12, 'tricyclist': 12, 'traffic_sign': 13, 'wheel': 14, 'plate': 15, 'face': 16, 'car_fake': 17, 'bicyclist_fake': 18, 'pedestrian_fake': 19, } DEFAULT_FACE_3D_CLASS_IDS = {0, 1, 2, 3, 4, 5, 6, 7, 8, 17} DEFAULT_COMPLETE_3D_CLASS_IDS = {9, 10, 11, 12, 18, 19} DEFAULT_TRACKED_CLASS_IDS = sorted(DEFAULT_FACE_3D_CLASS_IDS | DEFAULT_COMPLETE_3D_CLASS_IDS) VEHICLE_CLASS_NAME_ALIASES = { 'vehicle', 'car', 'car_fake', 'suv', 'pickup', 'medium_car', 'van', 'bus', 'truck', 'tanker', 'large_truck', 'construction_vehicle', 'special_vehicle', 'unknown', 'tricycle', 'tricyclist', } PEDESTRIAN_CLASS_NAME_ALIASES = { 'pedestrian', 'pedestrian_fake', 'person', 'ped', 'bicyclist', 'bicyclist_fake', 'motorcyclist', 'cyclist', 'rider', } DEFAULT_SUB_CLS = -1 DEFAULT_TIMESTAMP = 0 DEFAULT_MODEL_VERSION = '20260317' TIMESTAMP_CSV_NAME = 'frame_timestamps_interp.csv' TIMESTAMP_FRAME_ID_COLUMN = 'camera4_frame_id' TIMESTAMP_VALUE_COLUMN = 'camera4' DEFAULT_DUPLICATE_OVERLAP_THRESHOLD = 0.95 VRU_ASSOCIATION_CLASS_ID = -100 VRU_ASSOCIATION_TYPE_NAME = 'vru' def parse_numeric_value(value): """Convert CSV / JSON scalar strings to int or float when possible.""" if value is None: return None if isinstance(value, (int, float)): return value value_str = str(value).strip() if not value_str: return None # Preserve exact integer precision for large ids / timestamps such as # nanosecond log times written into exported filenames. if re.fullmatch(r'[+-]?\d+', value_str): try: return int(value_str) except ValueError: return value try: numeric = float(value_str) except ValueError: return value if numeric.is_integer(): return int(numeric) return numeric def normalize_class_name(name): """Normalize class-like strings to a lowercase underscore form.""" normalized = re.sub(r'[^a-z0-9]+', '_', str(name or '').strip().lower()) return normalized.strip('_') def invert_class_map(class_map): """Convert class_map{name->id} into the first canonical name per class id.""" class_names = {} for raw_name, class_id in class_map.items(): class_names.setdefault(int(class_id), normalize_class_name(raw_name)) return dict(sorted(class_names.items())) def is_vehicle_name(normalized_name): """Check whether a normalized class name belongs to the vehicle attribute task.""" return ( normalized_name in VEHICLE_CLASS_NAME_ALIASES or normalized_name == 'unknown' or 'vehicle' in normalized_name or 'truck' in normalized_name or 'bus' in normalized_name or 'tanker' in normalized_name or 'pickup' in normalized_name or normalized_name in {'car', 'suv', 'van', 'tricycle'} ) def is_pedestrian_name(normalized_name): """Check whether a normalized class name belongs to the pedestrian attribute task.""" return ( normalized_name in PEDESTRIAN_CLASS_NAME_ALIASES or 'pedestrian' in normalized_name or 'bicyclist' in normalized_name or 'motorcyclist' in normalized_name or 'cyclist' in normalized_name or 'rider' in normalized_name ) def resolve_task_class_ids(class_names): """Resolve vehicle / pedestrian class-id groups from canonical class names.""" vehicle_class_ids = set() pedestrian_class_ids = set() for cls_id, cls_name in class_names.items(): normalized_name = normalize_class_name(cls_name) if is_pedestrian_name(normalized_name): pedestrian_class_ids.add(int(cls_id)) continue if is_vehicle_name(normalized_name): vehicle_class_ids.add(int(cls_id)) return vehicle_class_ids, pedestrian_class_ids def load_tracking_metadata(data_config_path=DEFAULT_DATA_CONFIG_PATH, verbose=False): """Load Ground3D class metadata, falling back to synced in-script defaults.""" class_map = DEFAULT_CLASS_MAP.copy() face_3d_classes = set(DEFAULT_FACE_3D_CLASS_IDS) complete_3d_classes = set(DEFAULT_COMPLETE_3D_CLASS_IDS) source = 'builtin defaults' config_path = Path(data_config_path) if data_config_path else None if config_path and config_path.is_file(): if yaml is None: if verbose: print(f"Warning: PyYAML not available, falling back to built-in class map instead of {config_path}") else: try: with open(config_path, 'r', encoding='utf-8') as f: config = yaml.safe_load(f) or {} except Exception as exc: if verbose: print(f"Warning: Failed to parse data config {config_path}: {exc}. Using built-in class map.") else: raw_class_map = config.get('class_map') if isinstance(raw_class_map, dict) and raw_class_map: class_map = {str(key): int(value) for key, value in raw_class_map.items()} raw_face_3d = config.get('face_3d_classes') if raw_face_3d: face_3d_classes = {int(value) for value in raw_face_3d} raw_complete_3d = config.get('complete_3d_classes') if raw_complete_3d: complete_3d_classes = {int(value) for value in raw_complete_3d} source = str(config_path) elif config_path and verbose: print(f"Warning: Data config not found: {config_path}. Using built-in class map.") class_names = invert_class_map(class_map) tracked_classes = sorted((face_3d_classes | complete_3d_classes) or set(class_names)) vehicle_class_ids, pedestrian_class_ids = resolve_task_class_ids(class_names) return { 'source': source, 'class_names': class_names, 'face_3d_classes': face_3d_classes, 'complete_3d_classes': complete_3d_classes, 'tracked_classes': tracked_classes, 'vehicle_class_ids': vehicle_class_ids, 'pedestrian_class_ids': pedestrian_class_ids, } TRACKING_METADATA = load_tracking_metadata(verbose=False) CLASS_NAME_LOOKUP = TRACKING_METADATA['class_names'] VEHICLE_CLASS_IDS = TRACKING_METADATA['vehicle_class_ids'] PEDESTRIAN_CLASS_IDS = TRACKING_METADATA['pedestrian_class_ids'] TRACKED_CLASS_IDS = TRACKING_METADATA['tracked_classes'] def safe_int(value): """Best-effort integer conversion.""" numeric = parse_numeric_value(value) if numeric is None: return None try: return int(numeric) except (TypeError, ValueError): return None def normalize_heading_debug(heading_debug): """Normalize per-detection heading debug data to numeric Python types.""" if not isinstance(heading_debug, dict): return None normalized = {} yaw_bin = safe_int(heading_debug.get('yaw_bin')) if yaw_bin is not None: normalized['yaw_bin'] = yaw_bin yaw_delta = parse_numeric_value(heading_debug.get('yaw_delta')) if isinstance(yaw_delta, (int, float)): normalized['yaw_delta'] = float(yaw_delta) rot_y_decoded = parse_numeric_value(heading_debug.get('rot_y_decoded')) if isinstance(rot_y_decoded, (int, float)): normalized['rot_y_decoded'] = float(rot_y_decoded) yaw_probs_raw = heading_debug.get('yaw_probs') if isinstance(yaw_probs_raw, (list, tuple)): yaw_probs = [] valid = True for value in yaw_probs_raw: numeric = parse_numeric_value(value) if not isinstance(numeric, (int, float)): valid = False break yaw_probs.append(float(numeric)) if valid: normalized['yaw_probs'] = yaw_probs return normalized or None def resolve_output_version(model_version=None, existing_version=None): """Resolve the version string written into tracking detections.""" if model_version is not None: return str(model_version) if existing_version is not None: return str(existing_version) return DEFAULT_MODEL_VERSION def resolve_detection_type_name(det): """Resolve the best available normalized type name for one detection.""" for key in ('type_name', 'cls_name', 'class_name'): value = det.get(key) normalized = normalize_class_name(value) if normalized: return normalized class_id = safe_int(det.get('class_id')) if class_id is None: return '' return normalize_class_name(CLASS_NAME_LOOKUP.get(class_id, '')) def extract_frame_metadata_from_image_name(image_name): """Extract frame_id and optional timestamp from an image / JSON stem. Preferred filename format: camera4__.json Exported case / clip filename format also supported: ...__.json e.g. G1M3_xxx___.json 019dxxxx-clip_uuid__.json Legacy fallback: use the last numeric token as frame_id and leave timestamp unset. """ if not image_name: return None, None normalized = normalize_image_name(Path(image_name).stem) # Prefer explicit frame_id + timestamp carried in the filename, e.g. # camera4_36607_1760795616410702.json match = re.search(r'(?:^|_)camera\d+_(\d+)_(\d+)(?:$|_)', normalized) if match is not None: return int(match.group(1)), parse_numeric_value(match.group(2)) # Exported case / clip files often end with `__`, while # earlier tokens may include scene ids, datetimes, UUIDs, or saved indices. # Split on `_` rather than matching arbitrary digits so UUID-internal digits # do not interfere with the parsed tail fields. parts = normalized.split('_') if len(parts) >= 3 and parts[-1].isdigit() and parts[-2].isdigit(): return int(parts[-2]), parse_numeric_value(parts[-1]) match = re.search(r'(\d+)(?!.*\d)', normalized) if match is None: return None, None return int(match.group(1)), None def extract_frame_id_from_image_name(image_name): """Extract frame_id from an image / JSON stem.""" frame_id, _ = extract_frame_metadata_from_image_name(image_name) return frame_id def resolve_frame_metadata(image_name=None, frame_info=None, timestamp_lookup=None): """Resolve frame_id / timestamp with frame_info overrides when available. Priority: 1) frame_info.original_frame_id / frame_info.frame_id / frame_info.frameId 2) frame_id parsed from image_name Timestamp priority: 1) frame_info.timestamp 2) timestamp parsed from image_name 3) timestamp_lookup[frame_id] """ parsed_frame_id, parsed_timestamp = extract_frame_metadata_from_image_name(image_name) frame_info_frame_id = None frame_info_timestamp = None if isinstance(frame_info, dict): for key in ('original_frame_id', 'frame_id', 'frameId'): frame_info_frame_id = safe_int(frame_info.get(key)) if frame_info_frame_id is not None: break frame_info_timestamp = parse_numeric_value(frame_info.get('timestamp')) resolved_frame_id = frame_info_frame_id if frame_info_frame_id is not None else parsed_frame_id if frame_info_timestamp is not None: resolved_timestamp = frame_info_timestamp elif parsed_timestamp is not None: resolved_timestamp = parsed_timestamp elif timestamp_lookup and resolved_frame_id in timestamp_lookup: resolved_timestamp = timestamp_lookup[resolved_frame_id] else: resolved_timestamp = None return resolved_frame_id, resolved_timestamp def find_timestamp_csv(search_path): """Find frame_timestamps_interp.csv near an input file or directory.""" current = Path(search_path) if current.is_file(): current = current.parent for candidate_dir in [current, *current.parents]: candidate = candidate_dir / TIMESTAMP_CSV_NAME if candidate.is_file(): return candidate return None def load_frame_timestamp_lookup(search_path, verbose=True): """Load a frame_id -> timestamp lookup from frame_timestamps_interp.csv.""" csv_path = find_timestamp_csv(search_path) if csv_path is None: if verbose: print(f"Warning: {TIMESTAMP_CSV_NAME} not found near {search_path}") return {} with open(csv_path, 'r', encoding='utf-8-sig', newline='') as f: reader = csv.DictReader(f) if reader.fieldnames is None: if verbose: print(f"Warning: Empty timestamp CSV: {csv_path}") return {} field_map = {name.strip(): name for name in reader.fieldnames if name is not None} frame_id_key = field_map.get(TIMESTAMP_FRAME_ID_COLUMN) timestamp_key = field_map.get(TIMESTAMP_VALUE_COLUMN) if frame_id_key is None or timestamp_key is None: if verbose: print( f"Warning: Missing required columns in {csv_path}. " f"Expected '{TIMESTAMP_FRAME_ID_COLUMN}' and '{TIMESTAMP_VALUE_COLUMN}'." ) return {} timestamp_lookup = {} for row in reader: frame_id = safe_int(row.get(frame_id_key)) if frame_id is None: continue timestamp_value = parse_numeric_value(row.get(timestamp_key)) if timestamp_value is None: continue timestamp_lookup[frame_id] = timestamp_value if verbose: print(f"Loaded {len(timestamp_lookup)} frame timestamp(s) from {csv_path}") return timestamp_lookup def categorize_detection(det): """Map a detection to a coarse attribute category.""" attribute = det.get('attribute') if isinstance(attribute, dict): task = normalize_class_name(attribute.get('task', '')) if task == 'vehicle': return 'vehicle' if task == 'pedestrian': return 'pedestrian' type_name = resolve_detection_type_name(det) if type_name and is_vehicle_name(type_name): return 'vehicle' if type_name and is_pedestrian_name(type_name): return 'pedestrian' class_id = safe_int(det.get('class_id')) if class_id in VEHICLE_CLASS_IDS: return 'vehicle' if class_id in PEDESTRIAN_CLASS_IDS: return 'pedestrian' return None def class_id_is_vru(class_id): """Return whether a class id belongs to the pedestrian/rider association group.""" if class_id in PEDESTRIAN_CLASS_IDS: return True cls_name = normalize_class_name(CLASS_NAME_LOOKUP.get(class_id, '')) return bool(cls_name and is_pedestrian_name(cls_name)) def resolve_detection_association_class_id(det, association_mode='class'): """Resolve the class key used only for temporal association.""" class_id = det.get('class_id') if association_mode != 'vru': return class_id association_group = normalize_class_name(det.get('association_group') or det.get('association_type_name')) if association_group == VRU_ASSOCIATION_TYPE_NAME: return VRU_ASSOCIATION_CLASS_ID type_name = resolve_detection_type_name(det) if type_name and is_pedestrian_name(type_name): return VRU_ASSOCIATION_CLASS_ID if class_id_is_vru(class_id): return VRU_ASSOCIATION_CLASS_ID return class_id def association_type_name_for_class(association_class_id): if association_class_id == VRU_ASSOCIATION_CLASS_ID: return VRU_ASSOCIATION_TYPE_NAME return CLASS_NAME_LOOKUP.get(association_class_id, str(association_class_id)) def resolve_tracker_class_ids(target_classes, association_mode='class'): """Map target classes to tracker buckets without changing output class ids.""" tracker_class_ids = [] seen = set() for cls_id in target_classes: assoc_id = VRU_ASSOCIATION_CLASS_ID if association_mode == 'vru' and class_id_is_vru(cls_id) else cls_id if assoc_id in seen: continue tracker_class_ids.append(assoc_id) seen.add(assoc_id) return tracker_class_ids def compute_sub_cls(det): """Derive sub_cls from attribute outputs for vehicles and pedestrians.""" category = categorize_detection(det) attribute = det.get('attribute') if category is None or not isinstance(attribute, dict): return DEFAULT_SUB_CLS attr_cls = safe_int(attribute.get('attr_cls')) if attr_cls is None: return DEFAULT_SUB_CLS is_fake = safe_int(attribute.get('is_fake')) or 0 if category == 'vehicle': if is_fake == 1: return 26 if attr_cls <= 11: return attr_cls if attr_cls == 23: return 12 return attr_cls + 3 if category == 'pedestrian': return attr_cls return DEFAULT_SUB_CLS def enrich_detection_metadata(detection, image_name=None, timestamp_lookup=None, model_version=None, frame_info=None): """Inject tracking metadata and normalize optional heading debug info.""" type_name = resolve_detection_type_name(detection) if type_name: detection['type_name'] = type_name detection['version'] = resolve_output_version( model_version=model_version, existing_version=detection.get('version'), ) detection['sub_cls'] = compute_sub_cls(detection) frame_id, resolved_timestamp = resolve_frame_metadata( image_name=image_name, frame_info=frame_info, timestamp_lookup=timestamp_lookup, ) if frame_id is not None: detection['frameId'] = str(frame_id) else: detection.setdefault('frameId', None) if resolved_timestamp is not None: detection['timestamp'] = resolved_timestamp elif detection.get('timestamp') is None: detection['timestamp'] = DEFAULT_TIMESTAMP else: detection.setdefault('timestamp', DEFAULT_TIMESTAMP) heading_debug = normalize_heading_debug(detection.get('heading_debug')) if heading_debug is not None: detection['heading_debug'] = heading_debug elif 'heading_debug' in detection: detection.pop('heading_debug', None) return detection def enrich_predictions_data(predictions_data, timestamp_lookup=None, model_version=None): """Inject sub_cls / timestamp fields into existing frame predictions.""" enriched = [] for frame_data in predictions_data: image_name = frame_data.get('image_name') frame_info = frame_data.get('frame_info') detections = frame_data.get('detections', []) for det in detections: enrich_detection_metadata( det, image_name=image_name, timestamp_lookup=timestamp_lookup, model_version=model_version, frame_info=frame_info, ) enriched.append(frame_data) return enriched def limit_predictions_data(predictions_data, max_frames=None, verbose=True, source_name='input'): """Optionally keep only the first N frames in temporal order.""" if max_frames is None: return predictions_data if max_frames <= 0: return [] if len(predictions_data) <= max_frames: return predictions_data if verbose: print( f"Limiting {source_name} from {len(predictions_data)} frame(s) " f"to first {max_frames} frame(s)" ) return predictions_data[:max_frames] def compute_iou(box1, box2): """Compute IoU between two bounding boxes. Args: box1: [x1, y1, x2, y2] box2: [x1, y1, x2, y2] Returns: IoU value (0-1) """ x1_1, y1_1, x2_1, y2_1 = box1 x1_2, y1_2, x2_2, y2_2 = box2 # Compute intersection area x1_inter = max(x1_1, x1_2) y1_inter = max(y1_1, y1_2) x2_inter = min(x2_1, x2_2) y2_inter = min(y2_1, y2_2) if x2_inter < x1_inter or y2_inter < y1_inter: return 0.0 inter_area = (x2_inter - x1_inter) * (y2_inter - y1_inter) # Compute union area box1_area = (x2_1 - x1_1) * (y2_1 - y1_1) box2_area = (x2_2 - x1_2) * (y2_2 - y1_2) union_area = box1_area + box2_area - inter_area if union_area == 0: return 0.0 return inter_area / union_area def compute_intersection_over_min_area(box1, box2): """Measure how fully the smaller box is covered by the larger one.""" x1_1, y1_1, x2_1, y2_1 = box1 x1_2, y1_2, x2_2, y2_2 = box2 x1_inter = max(x1_1, x1_2) y1_inter = max(y1_1, y1_2) x2_inter = min(x2_1, x2_2) y2_inter = min(y2_1, y2_2) if x2_inter < x1_inter or y2_inter < y1_inter: return 0.0 inter_area = (x2_inter - x1_inter) * (y2_inter - y1_inter) box1_area = max(0.0, (x2_1 - x1_1) * (y2_1 - y1_1)) box2_area = max(0.0, (x2_2 - x1_2) * (y2_2 - y1_2)) min_area = min(box1_area, box2_area) if min_area <= 0.0: return 0.0 return inter_area / min_area def compute_distance(box1, box2): """Compute center distance between two bounding boxes. Args: box1: [x1, y1, x2, y2] box2: [x1, y1, x2, y2] Returns: Euclidean distance between centers """ cx1 = (box1[0] + box1[2]) / 2 cy1 = (box1[1] + box1[3]) / 2 cx2 = (box2[0] + box2[2]) / 2 cy2 = (box2[1] + box2[3]) / 2 return np.sqrt((cx1 - cx2)**2 + (cy1 - cy2)**2) def get_detection_confidence(det): """Return a sortable confidence score for one detection.""" confidence = parse_numeric_value(det.get('confidence', det.get('score'))) if isinstance(confidence, (int, float)): return float(confidence) return 0.0 def suppress_near_duplicate_detections( detections, overlap_threshold=DEFAULT_DUPLICATE_OVERLAP_THRESHOLD, class_key='class_id', ): """Remove same-association detections that almost completely overlap a higher-score box.""" if overlap_threshold <= 0 or len(detections) < 2: return list(detections), 0 ranked = sorted( enumerate(detections), key=lambda item: (-get_detection_confidence(item[1]), item[0]), ) kept = [] suppressed_count = 0 for det_idx, det in ranked: bbox = det.get('bbox') class_id = det.get(class_key, det.get('class_id')) if not isinstance(bbox, (list, tuple)) or len(bbox) < 4: kept.append((det_idx, det)) continue is_duplicate = False for _, kept_det in kept: if kept_det.get(class_key, kept_det.get('class_id')) != class_id: continue kept_bbox = kept_det.get('bbox') if not isinstance(kept_bbox, (list, tuple)) or len(kept_bbox) < 4: continue overlap = compute_intersection_over_min_area(kept_bbox, bbox) if overlap >= overlap_threshold: is_duplicate = True break if is_duplicate: suppressed_count += 1 continue kept.append((det_idx, det)) kept_indices = sorted(det_idx for det_idx, _ in kept) return [detections[det_idx] for det_idx in kept_indices], suppressed_count def compute_3d_distance(det1, det2): """Compute 3D spatial distance between two detections. Args: det1, det2: Detection dicts with 'object_3d' containing 'location' [x, y, z] Returns: 3D Euclidean distance in meters (or None if 3D info not available) """ if 'object_3d' not in det1 or 'object_3d' not in det2: return None obj1 = det1['object_3d'] obj2 = det2['object_3d'] # Support both dict and list formats if isinstance(obj1, dict) and isinstance(obj2, dict): if 'location' not in obj1 or 'location' not in obj2: return None loc1 = np.array(obj1['location'][:3]) # [x, y, z] loc2 = np.array(obj2['location'][:3]) elif isinstance(obj1, (list, tuple)) and isinstance(obj2, (list, tuple)): if len(obj1) < 3 or len(obj2) < 3: return None loc1 = np.array(obj1[:3]) # First 3 elements are [x, y, z] loc2 = np.array(obj2[:3]) else: return None return np.linalg.norm(loc1 - loc2) def compute_size_similarity(det1, det2): """Compute size similarity between two detections based on 3D dimensions. Args: det1, det2: Detection dicts with 'object_3d' containing 'dimensions' [l, h, w] Returns: Similarity score (0-1), higher is more similar (or None if 3D info not available) """ if 'object_3d' not in det1 or 'object_3d' not in det2: return None obj1 = det1['object_3d'] obj2 = det2['object_3d'] # Support both dict and list formats if isinstance(obj1, dict) and isinstance(obj2, dict): if 'dimensions' not in obj1 or 'dimensions' not in obj2: return None dim1 = np.array(obj1['dimensions'][:3]) # [l, h, w] dim2 = np.array(obj2['dimensions'][:3]) elif isinstance(obj1, (list, tuple)) and isinstance(obj2, (list, tuple)): if len(obj1) < 6 or len(obj2) < 6: return None dim1 = np.array(obj1[3:6]) # Elements [3,4,5] are [l, h, w] dim2 = np.array(obj2[3:6]) else: return None # Compute relative difference for each dimension diff = np.abs(dim1 - dim2) / (np.maximum(dim1, dim2) + 1e-6) similarity = 1.0 - np.mean(diff) return max(0.0, similarity) class SimpleTracker: """IoU-based object tracker with optional 3D distance matching.""" def __init__(self, iou_threshold=0.3, max_age=5, min_hits=1, distance_threshold=100, use_3d=False, max_3d_distance=10.0, w_iou=1.0, w_2d_dist=0.5, w_3d_dist=1.0, w_size=0.3, shared_id_counter=None): """Initialize tracker. Args: iou_threshold: Minimum IoU to match detections with tracks max_age: Maximum frames to keep track without detection min_hits: Minimum detections before track is confirmed distance_threshold: Maximum center distance for matching (pixels) use_3d: Whether to use 3D distance for matching max_3d_distance: Maximum 3D distance for matching (meters) w_iou: Weight for IoU in cost function w_2d_dist: Weight for 2D distance in cost function w_3d_dist: Weight for 3D distance in cost function (if use_3d=True) w_size: Weight for size similarity in cost function (if use_3d=True) shared_id_counter: Optional mutable list [next_id] shared across trackers to ensure globally unique track IDs across classes. """ self.iou_threshold = iou_threshold self.max_age = max_age self.min_hits = min_hits self.distance_threshold = distance_threshold self.use_3d = use_3d self.max_3d_distance = max_3d_distance self.w_iou = w_iou self.w_2d_dist = w_2d_dist self.w_3d_dist = w_3d_dist self.w_size = w_size # Use shared counter if provided, otherwise use a private one self._shared_id_counter = shared_id_counter if shared_id_counter is not None else [1] self.tracks = {} # track_id -> track_info def update(self, detections): """Update tracks with new detections. Args: detections: List of detection dicts with 'bbox', 'confidence', 'class_id', etc. Returns: List of detections with added 'track_id' field """ if not detections: # Age all tracks for track_id in list(self.tracks.keys()): self.tracks[track_id]['age'] += 1 if self.tracks[track_id]['age'] > self.max_age: del self.tracks[track_id] return [] # Match detections to existing tracks matched_detections = [] unmatched_detections = list(range(len(detections))) unmatched_tracks = list(self.tracks.keys()) # Build cost matrix (IoU + distance) if unmatched_tracks and unmatched_detections: cost_matrix = np.zeros((len(unmatched_tracks), len(unmatched_detections))) for i, track_id in enumerate(unmatched_tracks): track_box = self.tracks[track_id]['bbox'] track_class = self.tracks[track_id].get('association_class_id', self.tracks[track_id]['class_id']) track_obj3d = self.tracks[track_id].get('object_3d', None) for j, det_idx in enumerate(unmatched_detections): det_box = detections[det_idx]['bbox'] det_class = detections[det_idx].get('association_class_id', detections[det_idx]['class_id']) # Only match same class if track_class != det_class: cost_matrix[i, j] = 0.0 continue # Compute 2D IoU iou = compute_iou(track_box, det_box) # Compute 2D center distance distance_2d = compute_distance(track_box, det_box) # Initialize cost components cost = 0.0 # IoU component (primary) if iou > self.iou_threshold: cost += self.w_iou * iou else: cost_matrix[i, j] = 0.0 continue # 2D distance component if distance_2d < self.distance_threshold: cost += self.w_2d_dist * (1.0 - distance_2d / self.distance_threshold) else: cost_matrix[i, j] = 0.0 continue # 3D components (if enabled and available) if self.use_3d and track_obj3d is not None: # Create temporary track detection dict for 3D functions track_det = {'object_3d': track_obj3d} det_dict = detections[det_idx] # 3D distance component dist_3d = compute_3d_distance(track_det, det_dict) if dist_3d is not None and dist_3d < self.max_3d_distance: cost += self.w_3d_dist * (1.0 - dist_3d / self.max_3d_distance) elif dist_3d is not None and dist_3d >= self.max_3d_distance: # 3D distance too large, reject match cost_matrix[i, j] = 0.0 continue # Size similarity component size_sim = compute_size_similarity(track_det, det_dict) if size_sim is not None: cost += self.w_size * size_sim cost_matrix[i, j] = cost # Hungarian algorithm matching (optimal assignment) if SCIPY_AVAILABLE and cost_matrix.max() > 0: # Use Hungarian algorithm for optimal matching row_ind, col_ind = linear_sum_assignment(-cost_matrix) # Maximize (negative for minimize) matches = [(unmatched_tracks[i], unmatched_detections[j]) for i, j in zip(row_ind, col_ind) if cost_matrix[i, j] > 0] else: # Fallback to greedy matching if scipy not available matches = [] cost_copy = cost_matrix.copy() while True: max_val = cost_copy.max() if max_val == 0: break i, j = np.unravel_index(cost_copy.argmax(), cost_copy.shape) track_id = unmatched_tracks[i] det_idx = unmatched_detections[j] matches.append((track_id, det_idx)) # Mark as matched cost_copy[i, :] = 0 cost_copy[:, j] = 0 # Update matched tracks for track_id, det_idx in matches: self.tracks[track_id]['bbox'] = detections[det_idx]['bbox'] self.tracks[track_id]['age'] = 0 self.tracks[track_id]['hits'] += 1 # Update 3D info if available if 'object_3d' in detections[det_idx]: self.tracks[track_id]['object_3d'] = detections[det_idx]['object_3d'] # Add track_id to detection detections[det_idx]['track_id'] = track_id matched_detections.append(det_idx) unmatched_tracks.remove(track_id) unmatched_detections.remove(det_idx) # Age unmatched tracks for track_id in unmatched_tracks: self.tracks[track_id]['age'] += 1 if self.tracks[track_id]['age'] > self.max_age: del self.tracks[track_id] # Create new tracks for unmatched detections for det_idx in unmatched_detections: track_id = self._shared_id_counter[0] self._shared_id_counter[0] += 1 track_info = { 'bbox': detections[det_idx]['bbox'], 'class_id': detections[det_idx]['class_id'], 'association_class_id': detections[det_idx].get('association_class_id', detections[det_idx]['class_id']), 'age': 0, 'hits': 1 } # Store 3D info if available if 'object_3d' in detections[det_idx]: track_info['object_3d'] = detections[det_idx]['object_3d'] self.tracks[track_id] = track_info # Add track_id to detection detections[det_idx]['track_id'] = track_id matched_detections.append(det_idx) # Return detections with track IDs (only confirmed tracks) result = [] for det_idx in matched_detections: det = detections[det_idx] track_id = det['track_id'] # Only include if track has enough hits if self.tracks[track_id]['hits'] >= self.min_hits: result.append(det) return result def parse_det_format(det_dict, image_name=None, timestamp_lookup=None, model_version=None, frame_info=None): """Parse a single-frame detection result in det_format into internal frame data format. Args: det_dict: Either a flat dict of detections keyed by string index (as in det_format.json reference), OR a dict with a top-level "detections" key wrapping that flat dict (actual merge_json format). Each detection entry has: type, score, roi_id, box2d, xyzlhwyaw, face_cls, cut_cls. type_name and cut_cls_name are optional. image_name: Optional frame name / image stem. frame_info: Optional frame metadata dict. When it contains original_frame_id/frame_id/frameId, that value is preferred over any frame id parsed from image_name. Returns: Frame data dict with keys 'image_name' and 'detections', compatible with the list expected by track_objects(). """ # Support both flat format and wrapped {"detections": {...}} format if 'detections' in det_dict and isinstance(det_dict['detections'], dict): raw_detections = det_dict['detections'] else: raw_detections = det_dict face_map = { 'front': 'kMonocular3DFront', 'tail': 'kMonocular3DRear', 'back': 'kMonocular3DRear', 'left': 'kMonocular3DLeft', 'right': 'kMonocular3DRight', 'center': 'kMonocular3DCenter', 'none': 'kMonocular3DCenter' } resolved_frame_id, resolved_timestamp = resolve_frame_metadata( image_name=image_name, frame_info=frame_info, timestamp_lookup=timestamp_lookup, ) detections = [] for det in raw_detections.values(): class_id = int(det['type']) bbox = [float(v) for v in det['box2d']] score = float(det['score']) # Parse 3D info from xyzlhwyaw; sentinel value -1 means no 3D xyzlhwyaw_raw = det.get('xyzlhwyaw', []) object_3d = None if xyzlhwyaw_raw and float(xyzlhwyaw_raw[0]) != -1: object_3d = [float(v) for v in xyzlhwyaw_raw] # [x, y, z, l, h, w, yaw] # Parse 3D info from xyzlhwyaw_ego; sentinel value -1 means no 3D xyzlhwyaw_ego_raw = det.get('xyzlhwyaw_ego', []) object_3d_ego = None if xyzlhwyaw_ego_raw and float(xyzlhwyaw_ego_raw[0]) != -1: object_3d_ego = [float(v) for v in xyzlhwyaw_ego_raw] # [x, y, z, l, h, w, yaw] # Parse 3D box center (separate from face-center xyzlhwyaw) box_center_xyz_raw = det.get('box_center_xyz', []) box_center_3d = None if box_center_xyz_raw and len(box_center_xyz_raw) >= 3: box_center_3d = [float(v) for v in box_center_xyz_raw[:3]] # [x, y, z] box_center_xyz_ego_raw = det.get('box_center_xyz_ego', []) box_center_3d_ego = None if box_center_xyz_ego_raw and len(box_center_xyz_ego_raw) >= 3: box_center_3d_ego = [float(v) for v in box_center_xyz_ego_raw[:3]] # [x, y, z] detection = { 'bbox': bbox, 'confidence': score, 'class_id': class_id, 'type_name': det.get('type_name', ''), 'face_cls': det.get('face_cls', 'none'), 'cut_cls': int(det.get('cut_cls', -1)), 'cut_cls_name': det.get('cut_cls_name', 'none'), 'frameId': str(resolved_frame_id) if resolved_frame_id is not None else None, 'version': resolve_output_version(model_version=model_version, existing_version=det.get('version')), 'timestamp': resolved_timestamp if resolved_timestamp is not None else DEFAULT_TIMESTAMP, 'roi_id': int(det.get('roi_id', -1)), } detection['anchor'] = face_map.get(detection['face_cls'], 'kMonocular3DCenter') if 'attribute' in det: detection['attribute'] = det.get('attribute') if 'heading_debug' in det: detection['heading_debug'] = det.get('heading_debug') for field_name in ( 'difficulty_logit', 'difficulty_prob', 'difficulty_label', 'difficulty_name', ): if field_name in det: detection[field_name] = det.get(field_name) for field_name in ( 'association_group', 'association_class_id', 'association_type_name', 'raw_cls_id', 'raw_cls_name', ): if field_name in det: detection[field_name] = det.get(field_name) if object_3d is not None: detection['object_3d'] = object_3d if box_center_3d is not None: detection['box_center_3d'] = box_center_3d if object_3d_ego is not None: detection['object_3d_ego'] = object_3d_ego if box_center_3d_ego is not None: detection['box_center_3d_ego'] = box_center_3d_ego enrich_detection_metadata( detection, image_name=image_name, timestamp_lookup=timestamp_lookup, model_version=model_version, frame_info=frame_info, ) detections.append(detection) return { 'image_name': image_name, 'detections': detections, } def load_predictions_from_dir(input_dir, pattern='*.json', verbose=True, timestamp_lookup=None, max_frames=None, model_version=None): """Load per-frame detection results from a directory of det_format JSON files. Files are sorted lexicographically (i.e. by filename) to preserve temporal order. Args: input_dir: Directory containing per-frame JSON files in det_format. pattern: Glob pattern used to match JSON files (default: '*.json'). Returns: List of frame data dicts compatible with track_objects(). """ input_dir = Path(input_dir) json_files = sorted(input_dir.glob(pattern)) if not json_files: if verbose: print(f"Warning: No files matching '{pattern}' found in {input_dir}") return [] if verbose: print(f"Found {len(json_files)} frame file(s) in {input_dir}") if max_frames is not None and max_frames > 0 and len(json_files) > max_frames: if verbose: print(f"Limiting {input_dir} to first {max_frames} frame file(s)") json_files = json_files[:max_frames] if timestamp_lookup is None: timestamp_lookup = load_frame_timestamp_lookup(input_dir, verbose=verbose) predictions_data = [] for json_file in json_files: with open(json_file, 'r', encoding='utf-8') as f: det_dict = json.load(f) frame_data = parse_det_format( det_dict, image_name=json_file.stem, timestamp_lookup=timestamp_lookup, model_version=model_version, ) predictions_data.append(frame_data) return predictions_data def track_objects(predictions_data, target_classes=None, iou_threshold=0.3, max_age=5, min_hits=1, distance_threshold=100, use_3d=False, max_3d_distance=10.0, verbose=True, duplicate_overlap_threshold=DEFAULT_DUPLICATE_OVERLAP_THRESHOLD, association_mode='class'): """Track objects across frames. Args: predictions_data: List of frame predictions from predictions.json target_classes: List of class IDs to track. Defaults to Ground3D face_3d_classes ∪ complete_3d_classes. iou_threshold: Minimum IoU for matching max_age: Maximum frames without detection min_hits: Minimum detections before confirmed distance_threshold: Maximum center distance for matching (pixels) use_3d: Whether to use 3D distance for matching max_3d_distance: Maximum 3D distance for matching (meters) duplicate_overlap_threshold: Suppress same-class detections when the smaller box is almost fully covered by a higher-score box. Returns: List of frame predictions with track IDs added """ if target_classes is None: target_classes = list(TRACKED_CLASS_IDS) target_classes = [int(cls_id) for cls_id in target_classes] target_class_set = set(target_classes) tracker_class_ids = resolve_tracker_class_ids(target_classes, association_mode=association_mode) # Shared counter ensures track IDs are unique across all classes shared_id_counter = [1] trackers = {cls_id: SimpleTracker(iou_threshold, max_age, min_hits, distance_threshold, use_3d=use_3d, max_3d_distance=max_3d_distance, shared_id_counter=shared_id_counter) for cls_id in tracker_class_ids} tracking_results = [] for frame_idx, frame_data in enumerate(predictions_data): if verbose: print(f"Processing frame {frame_idx + 1}/{len(predictions_data)}: {frame_data.get('image_name', 'unknown')}") # Group detections by class detections_by_class = {cls_id: [] for cls_id in tracker_class_ids} non_tracked_detections = [] for det in frame_data.get('detections', []): cls_id = safe_int(det.get('class_id')) if cls_id in target_class_set: det['class_id'] = cls_id association_class_id = resolve_detection_association_class_id(det, association_mode=association_mode) det['association_class_id'] = association_class_id det['association_type_name'] = association_type_name_for_class(association_class_id) if association_class_id == VRU_ASSOCIATION_CLASS_ID: det['association_group'] = VRU_ASSOCIATION_TYPE_NAME detections_by_class.setdefault(association_class_id, []).append(det) else: non_tracked_detections.append(det) # Track each class separately tracked_detections = [] frame_duplicate_count = 0 for cls_id in tracker_class_ids: if cls_id in detections_by_class: deduped_detections, suppressed_count = suppress_near_duplicate_detections( detections_by_class[cls_id], overlap_threshold=duplicate_overlap_threshold, class_key='association_class_id' if association_mode == 'vru' else 'class_id', ) frame_duplicate_count += suppressed_count tracked = trackers[cls_id].update(deduped_detections) tracked_detections.extend(tracked) if verbose and frame_duplicate_count > 0: print( f" Suppressed {frame_duplicate_count} near-duplicate detection(s) " f"before association" ) # Append non-tracked classes directly (no tracking needed) tracked_detections.extend(non_tracked_detections) # Create output frame data result_frame = { 'image_name': frame_data.get('image_name'), 'detections': tracked_detections } # Copy frame_info if present if 'frame_info' in frame_data: result_frame['frame_info'] = frame_data['frame_info'] # Add tracking statistics result_frame['tracking_stats'] = { 'total_tracks': sum(len(tracker.tracks) for tracker in trackers.values()), 'active_tracks_by_class': {cls_id: len([t for t in tracker.tracks.values() if t['age'] == 0]) for cls_id, tracker in trackers.items()} } tracking_results.append(result_frame) return tracking_results def save_tracking_results(tracking_results, output_path): """Save tracking results to a JSON file.""" output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) with open(output_path, 'w', encoding='utf-8') as f: json.dump(tracking_results, f, indent=2, ensure_ascii=False) def count_unique_tracks(tracking_results): """Count unique track ids in tracking output.""" total_tracks = set() for frame in tracking_results: for det in frame.get('detections', []): if 'track_id' in det: total_tracks.add(det['track_id']) return len(total_tracks) def run_tracking_job(input_dir, output_path, file_pattern, classes, iou_threshold, max_age, min_hits, distance_threshold, use_3d, max_3d_distance, verbose=True, timestamp_lookup=None, max_frames=None, model_version=None, association_mode='class'): """Run one tracking job from a directory of per-frame JSON files.""" predictions_data = load_predictions_from_dir( input_dir, file_pattern, verbose=verbose, timestamp_lookup=timestamp_lookup, max_frames=max_frames, model_version=model_version, ) if not predictions_data: return { 'ok': False, 'input_dir': str(input_dir), 'output_path': str(output_path), 'frames': 0, 'unique_tracks': 0, 'reason': 'no_frames_loaded', } tracking_results = track_objects( predictions_data, target_classes=classes, iou_threshold=iou_threshold, max_age=max_age, min_hits=min_hits, distance_threshold=distance_threshold, use_3d=use_3d, max_3d_distance=max_3d_distance, verbose=verbose, association_mode=association_mode, ) save_tracking_results(tracking_results, output_path) return { 'ok': True, 'input_dir': str(input_dir), 'output_path': str(output_path), 'frames': len(predictions_data), 'unique_tracks': count_unique_tracks(tracking_results), } def find_case_dirs(results_root, model_name): """Scan the model directory once and return case dirs containing tracking inputs.""" search_root = Path(results_root) / model_name if not search_root.exists(): print(f"Error: model directory does not exist: {search_root}") return [] case_dirs = set() target_dirs = {'roi0', 'roi1', 'merge_json'} for current_dir, dirnames, _ in os.walk(search_root): if target_dirs.intersection(dirnames): case_dirs.add(Path(current_dir)) return sorted(case_dirs) def process_case(case_dir, file_pattern, classes, iou_threshold, max_age, min_hits, distance_threshold, use_3d, max_3d_distance, merge_output_name, max_frames=None, model_version=None, association_mode='class'): """Track all sources for one case and write the merged output.""" case_dir = Path(case_dir) timestamp_lookup = load_frame_timestamp_lookup(case_dir, verbose=False) source_specs = [ ('roi0', 'roi0.json'), ('roi1', 'roi1.json'), ('merge_json', 'merge.json'), ] tracking_outputs = {} source_summaries = [] missing_sources = [] for input_name, output_name in source_specs: input_dir = case_dir / input_name output_path = case_dir / output_name if not input_dir.is_dir(): missing_sources.append(input_name) continue result = run_tracking_job( input_dir=input_dir, output_path=output_path, file_pattern=file_pattern, classes=classes, iou_threshold=iou_threshold, max_age=max_age, min_hits=min_hits, distance_threshold=distance_threshold, use_3d=use_3d, max_3d_distance=max_3d_distance, verbose=False, timestamp_lookup=timestamp_lookup, max_frames=max_frames, model_version=model_version, association_mode=association_mode, ) result['source'] = input_name source_summaries.append(result) if result['ok']: tracking_outputs[input_name] = output_path merge_ok = merge_case( roi0_path=tracking_outputs.get('roi0'), roi1_path=tracking_outputs.get('roi1'), merge_path=tracking_outputs.get('merge_json'), output_path=case_dir / merge_output_name, source_names=('roi0', 'roi1', 'merge'), verbose=False, ) ok_sources = [item['source'] for item in source_summaries if item['ok']] failed_sources = [item['source'] for item in source_summaries if not item['ok']] return { 'case_dir': str(case_dir), 'ok': bool(ok_sources) and merge_ok and not failed_sources, 'tracked_sources': ok_sources, 'failed_sources': failed_sources, 'missing_sources': missing_sources, 'merge_ok': merge_ok, 'source_summaries': source_summaries, 'merge_output': str(case_dir / merge_output_name), } def print_case_summary(result): """Print a clean per-case summary after batch processing.""" print("") print("==========================================") print(f"Case : {result['case_dir']}") print(f"Merge : {'ok' if result['merge_ok'] else 'failed'}") print(f"Output : {result['merge_output']}") if result['tracked_sources']: print(f"Tracked: {', '.join(result['tracked_sources'])}") if result['failed_sources']: print(f"Failed : {', '.join(result['failed_sources'])}") if result['missing_sources']: print(f"Missing: {', '.join(result['missing_sources'])}") for item in result['source_summaries']: status = 'ok' if item['ok'] else item.get('reason', 'failed') print( f" - {item['source']}: {status}, frames={item['frames']}, " f"tracks={item['unique_tracks']}" ) print("==========================================") def track_cases_in_batch(results_root, model_name, file_pattern, classes, iou_threshold, max_age, min_hits, distance_threshold, use_3d, max_3d_distance, merge_output_name, num_workers, max_frames=None, model_version=None, association_mode='class'): """Run tracking and merge end-to-end for every case under one model directory.""" case_dirs = find_case_dirs(results_root, model_name) if not case_dirs: print("Error: No case directories containing roi0/roi1/merge_json were found.") return False print(f"Found {len(case_dirs)} case(s) under {Path(results_root) / model_name}") print(f"Using {num_workers} worker(s)") results = [] if num_workers <= 1: for case_dir in case_dirs: result = process_case( case_dir=case_dir, file_pattern=file_pattern, classes=classes, iou_threshold=iou_threshold, max_age=max_age, min_hits=min_hits, distance_threshold=distance_threshold, use_3d=use_3d, max_3d_distance=max_3d_distance, merge_output_name=merge_output_name, max_frames=max_frames, model_version=model_version, association_mode=association_mode, ) results.append(result) print_case_summary(result) else: with ProcessPoolExecutor(max_workers=num_workers) as executor: future_to_case = { executor.submit( process_case, case_dir, file_pattern, classes, iou_threshold, max_age, min_hits, distance_threshold, use_3d, max_3d_distance, merge_output_name, max_frames, model_version, association_mode, ): case_dir for case_dir in case_dirs } for future in as_completed(future_to_case): result = future.result() results.append(result) print_case_summary(result) succeeded = sum(1 for item in results if item['ok']) failed = len(results) - succeeded print("") print("Batch tracking complete") print(f" succeeded_cases: {succeeded}") print(f" failed_cases: {failed}") return failed == 0 def main(): parser = argparse.ArgumentParser(description='Track objects from predictions.json') parser.add_argument('--input', type=str, default=None, help='Input predictions.json file (multi-frame list format) ' 'or a single-frame det_format JSON file') parser.add_argument('--input-dir', type=str, default=None, help='Directory containing per-frame JSON files in det_format ' '(used when inference results are saved as individual files)') parser.add_argument('--results-root', type=str, default=None, help='Batch mode root directory that contains per-model case results') parser.add_argument('--model-name', type=str, default=None, help='Batch mode model directory name under results-root') parser.add_argument('--file-pattern', type=str, default='*.json', help='Glob pattern for JSON files inside --input-dir (default: %(default)s)') parser.add_argument('--output', type=str, default=None, help='Output tracking.json file path (default: same dir as input)') parser.add_argument('--merge-output-name', type=str, default='combined_tracking.json', help='Batch mode output filename written inside each case directory') parser.add_argument('--num-workers', type=int, default=max(1, min(8, os.cpu_count() or 1)), help='Batch mode worker count for case-level parallelism (default: %(default)s)') parser.add_argument( '--classes', type=int, nargs='+', default=None, help='Class IDs to track (default: Ground3D face_3d_classes + complete_3d_classes)', ) parser.add_argument('--iou-threshold', type=float, default=0.3, help='IoU threshold for matching (default: 0.3)') parser.add_argument('--max-age', type=int, default=5, help='Maximum frames without detection (default: 5)') parser.add_argument('--min-hits', type=int, default=1, help='Minimum detections before confirmed (default: 1)') parser.add_argument('--distance-threshold', type=float, default=100, help='Maximum center distance for matching in pixels (default: 100)') parser.add_argument('--max-frames', type=int, default=None, help='Only process the first N frames in temporal order (default: all frames)') parser.add_argument('--model-version', type=str, default=None, help=f"Version string written into detections (default: keep existing or '{DEFAULT_MODEL_VERSION}' when missing)") parser.add_argument('--use-3d', action='store_true', help='Use 3D distance for matching (requires object_3d in detections)') parser.add_argument('--max-3d-distance', type=float, default=10.0, help='Maximum 3D distance for matching in meters (default: 10.0)') parser.add_argument( '--association-mode', type=str, choices=('class', 'vru'), default='class', help='Temporal association class mapping. "vru" tracks pedestrian/rider classes as one VRU group while preserving raw classes.', ) args = parser.parse_args() if args.classes is None: args.classes = list(TRACKED_CLASS_IDS) if args.max_frames is not None and args.max_frames <= 0: parser.error('--max-frames must be a positive integer') if args.results_root or args.model_name: if not args.results_root or not args.model_name: parser.error('--results-root and --model-name must be provided together in batch mode') track_cases_in_batch( results_root=args.results_root, model_name=args.model_name, file_pattern=args.file_pattern, classes=args.classes, iou_threshold=args.iou_threshold, max_age=args.max_age, min_hits=args.min_hits, distance_threshold=args.distance_threshold, use_3d=args.use_3d, max_3d_distance=args.max_3d_distance, merge_output_name=args.merge_output_name, num_workers=args.num_workers, max_frames=args.max_frames, model_version=args.model_version, association_mode=args.association_mode, ) return # Read input predictions if args.input_dir is not None: # New format: directory of per-frame det_format JSON files predictions_data = load_predictions_from_dir( args.input_dir, args.file_pattern, max_frames=args.max_frames, model_version=args.model_version, ) if not predictions_data: print("Error: No frame files loaded. Check --input-dir and --file-pattern.") return input_path = Path(args.input_dir) # used for default output path elif args.input is not None: input_path = Path(args.input) if not input_path.exists(): print(f"Error: Input file not found: {input_path}") return print(f"Loading predictions from {input_path}") timestamp_lookup = load_frame_timestamp_lookup(input_path, verbose=False) with open(input_path, 'r', encoding='utf-8') as f: raw_data = json.load(f) # Auto-detect format: dict → single-frame det_format; list → multi-frame predictions.json if isinstance(raw_data, dict): print("Detected single-frame det_format JSON. Wrapping into one-frame sequence.") predictions_data = [ parse_det_format( raw_data, image_name=input_path.stem, timestamp_lookup=timestamp_lookup, model_version=args.model_version, ) ] else: predictions_data = raw_data predictions_data = enrich_predictions_data( predictions_data, timestamp_lookup=timestamp_lookup, model_version=args.model_version, ) predictions_data = limit_predictions_data( predictions_data, max_frames=args.max_frames, source_name=str(input_path), ) else: print("Error: Either --input or --input-dir must be specified.") return print(f"Loaded {len(predictions_data)} frame(s)") # Perform tracking print(f"Tracking objects of classes: {args.classes}") resolved_class_names = [CLASS_NAME_LOOKUP.get(int(cls_id), str(cls_id)) for cls_id in args.classes] print(f"Tracking class names: {', '.join(resolved_class_names)}") print(f"Tracking class config source: {TRACKING_METADATA['source']}") if args.max_frames is not None: print(f"Tracking only first {args.max_frames} frame(s)") if args.model_version is not None: print(f"Tracking model version: {args.model_version}") if args.use_3d: print(f"Using 3D distance matching (max 3D distance: {args.max_3d_distance}m)") print(f"Association mode: {args.association_mode}") tracking_results = track_objects( predictions_data, target_classes=args.classes, iou_threshold=args.iou_threshold, max_age=args.max_age, min_hits=args.min_hits, distance_threshold=args.distance_threshold, use_3d=args.use_3d, max_3d_distance=args.max_3d_distance, association_mode=args.association_mode, ) # Determine output path if args.output is None: out_base = input_path if input_path.is_dir() else input_path.parent output_path = out_base / 'tracking.json' else: output_path = Path(args.output) # Save results print(f"Saving tracking results to {output_path}") save_tracking_results(tracking_results, output_path) # Print statistics total_track_count = count_unique_tracks(tracking_results) print(f"\nTracking completed!") print(f"Total unique tracks: {total_track_count}") print(f"Output saved to: {output_path}") if __name__ == '__main__': main()