1709 lines
64 KiB
Python
Executable File
1709 lines
64 KiB
Python
Executable File
"""
|
||
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_<frame_id>_<timestamp>.json
|
||
|
||
Exported case / clip filename format also supported:
|
||
..._<frame_id>_<timestamp>.json
|
||
e.g. G1M3_xxx_<saved_idx>_<frame_id>_<timestamp>.json
|
||
019dxxxx-clip_uuid_<frame_id>_<timestamp>.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 `_<frame_id>_<timestamp>`, 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()
|