zhao.zhu/yolov26_3d
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
yolov26_3d/tools/temporal_analysis/track_objects.py

1709 lines
64 KiB
Python
Raw Permalink Normal View History

单目3D初始代码
2026-06-24 09:35:46 +08:00
"""
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()
Reference in New Issue Copy Permalink