367 lines
13 KiB
Python
367 lines
13 KiB
Python
|
"""
|
|||
|
|
Merge multiple per-source tracking JSON files into a single combined JSON.
|
|||
|
|
|
|||
|
|
Each source (roi0, roi1, merge) is the output of track_objects.py. The script
|
|||
|
|
aligns frames by image_name, adds a "source" field to every detection, and
|
|||
|
|
writes a unified JSON where each frame contains detections from all sources.
|
|||
|
|
|
|||
|
|
Output format (list of frame dicts):
|
|||
|
|
[
|
|||
|
|
{
|
|||
|
|
"image_name": "1741824631_0",
|
|||
|
|
"detections": [
|
|||
|
|
{ "lane_assignment": 0, "track_id": 1, "class_id": 0, "bbox": [...], ... },
|
|||
|
|
{ "lane_assignment": 1, "track_id": 1, ... },
|
|||
|
|
{ "lane_assignment": 2, "track_id": 2, ... },
|
|||
|
|
...
|
|||
|
|
],
|
|||
|
|
"tracking_stats": {
|
|||
|
|
"roi0": { "total_tracks": 3, "active_tracks_by_class": {...} },
|
|||
|
|
"roi1": { ... },
|
|||
|
|
"merge": { ... }
|
|||
|
|
}
|
|||
|
|
},
|
|||
|
|
...
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
Notes:
|
|||
|
|
- track_id values are offset per source so they are globally unique:
|
|||
|
|
roi0 keeps original IDs (0–9999), roi1 adds 10000, merge adds 20000.
|
|||
|
|
- lane_assignment values: 0=roi0, 1=roi1, 2=merge (positional index in source_map).
|
|||
|
|
- Missing sources are silently skipped (partial merge is valid).
|
|||
|
|
- Frames present in some sources but not others are still included, with
|
|||
|
|
detections only from the sources that have that frame.
|
|||
|
|
|
|||
|
|
Usage:
|
|||
|
|
python merge_tracking_results.py \
|
|||
|
|
--roi0 {case_dir}/roi0.json \
|
|||
|
|
--roi1 {case_dir}/roi1.json \
|
|||
|
|
--merge {case_dir}/merge.json \
|
|||
|
|
--output {case_dir}/combined_tracking.json
|
|||
|
|
|
|||
|
|
# Batch mode: scan all cases under results_root/model_name
|
|||
|
|
python merge_tracking_results.py \
|
|||
|
|
--results-root /data1/dongying/Mono3d/G1M3/cases_regular \
|
|||
|
|
--model-name model_20260228
|
|||
|
|
|
|||
|
|
# Custom source names (optional):
|
|||
|
|
python merge_tracking_results.py \
|
|||
|
|
--roi0 roi0.json --roi1 roi1.json --merge merge.json \
|
|||
|
|
--source-names roi0_track roi1_track merge_track \
|
|||
|
|
--output combined_tracking.json
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import json
|
|||
|
|
import re
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def normalize_image_name(name):
|
|||
|
|
"""Normalize image_name to a canonical form for cross-source frame alignment.
|
|||
|
|
|
|||
|
|
merge_json filenames are saved as ``{saved_count:06d}_{img_name}_merged.json``
|
|||
|
|
while roi0/roi1 use plain ``{img_name}.json``. This function strips both
|
|||
|
|
the leading 6-digit counter prefix and the trailing ``_merged`` suffix so
|
|||
|
|
all three sources resolve to the same key.
|
|||
|
|
|
|||
|
|
Examples::
|
|||
|
|
|
|||
|
|
'000001_1741824631_0_merged' -> '1741824631_0'
|
|||
|
|
'1741824631_0' -> '1741824631_0' (unchanged)
|
|||
|
|
"""
|
|||
|
|
name = re.sub(r'^\d{6}_', '', name) # strip leading counter, e.g. '000001_'
|
|||
|
|
name = re.sub(r'_merged$', '', name) # strip trailing '_merged'
|
|||
|
|
return name
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_tracking_json(path, verbose=True):
|
|||
|
|
"""Load a tracking JSON file produced by track_objects.py.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
path: Path to the JSON file (list of frame dicts).
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List of frame dicts, or empty list if file does not exist / is invalid.
|
|||
|
|
"""
|
|||
|
|
p = Path(path)
|
|||
|
|
if not p.exists():
|
|||
|
|
if verbose:
|
|||
|
|
print(f"Warning: tracking file not found, skipping: {p}")
|
|||
|
|
return []
|
|||
|
|
with open(p, 'r', encoding='utf-8') as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
if not isinstance(data, list):
|
|||
|
|
if verbose:
|
|||
|
|
print(f"Warning: expected a list in {p}, got {type(data).__name__}. Skipping.")
|
|||
|
|
return []
|
|||
|
|
if verbose:
|
|||
|
|
print(f"Loaded {len(data)} frame(s) from {p}")
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_source_map(roi0_path, roi1_path, merge_path, source_names, extra_sources=None, verbose=True):
|
|||
|
|
"""Load all available tracking inputs for a single case."""
|
|||
|
|
roi0_name, roi1_name, merge_name = source_names
|
|||
|
|
|
|||
|
|
source_map = []
|
|||
|
|
if roi0_path is not None:
|
|||
|
|
frames = load_tracking_json(roi0_path, verbose=verbose)
|
|||
|
|
if frames:
|
|||
|
|
source_map.append((roi0_name, frames))
|
|||
|
|
if roi1_path is not None:
|
|||
|
|
frames = load_tracking_json(roi1_path, verbose=verbose)
|
|||
|
|
if frames:
|
|||
|
|
source_map.append((roi1_name, frames))
|
|||
|
|
if merge_path is not None:
|
|||
|
|
frames = load_tracking_json(merge_path, verbose=verbose)
|
|||
|
|
if frames:
|
|||
|
|
source_map.append((merge_name, frames))
|
|||
|
|
for source_name, source_path in extra_sources or []:
|
|||
|
|
frames = load_tracking_json(source_path, verbose=verbose)
|
|||
|
|
if frames:
|
|||
|
|
source_map.append((source_name, frames))
|
|||
|
|
|
|||
|
|
return source_map
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Offset added to track_id per source so IDs are globally unique:
|
|||
|
|
# source_idx 0 (roi0) -> no offset
|
|||
|
|
# source_idx 1 (roi1) -> +10000
|
|||
|
|
# source_idx 2 (merge) -> +20000
|
|||
|
|
# source_idx 3+ (extra src) -> keep incrementing by +10000
|
|||
|
|
TRACK_ID_OFFSET_PER_SOURCE = 10000
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_tracking_results(source_map):
|
|||
|
|
"""Merge tracking results from multiple sources into a single per-frame list.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
source_map: Ordered dict / list of (source_name, frames_list) pairs.
|
|||
|
|
Each frames_list is the output of track_objects.py.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List of merged frame dicts, one per unique image_name found across all
|
|||
|
|
sources. Frames are ordered by the first source that contains them;
|
|||
|
|
remaining sources are appended in order.
|
|||
|
|
"""
|
|||
|
|
# Build per-source index: image_name -> frame_dict
|
|||
|
|
source_indices = {}
|
|||
|
|
all_image_names_ordered = []
|
|||
|
|
seen_names = set()
|
|||
|
|
|
|||
|
|
for source_name, frames in source_map:
|
|||
|
|
idx = {}
|
|||
|
|
for frame in frames:
|
|||
|
|
name = normalize_image_name(frame.get('image_name') or '')
|
|||
|
|
idx[name] = frame
|
|||
|
|
if name not in seen_names:
|
|||
|
|
all_image_names_ordered.append(name)
|
|||
|
|
seen_names.add(name)
|
|||
|
|
source_indices[source_name] = idx
|
|||
|
|
|
|||
|
|
merged_frames = []
|
|||
|
|
for image_name in all_image_names_ordered:
|
|||
|
|
merged_detections = []
|
|||
|
|
merged_stats = {}
|
|||
|
|
|
|||
|
|
for source_idx, (source_name, _) in enumerate(source_map):
|
|||
|
|
frame = source_indices[source_name].get(image_name)
|
|||
|
|
if frame is None:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# Tag every detection with its lane_assignment (0=roi0, 1=roi1, 2=merge)
|
|||
|
|
# and offset track_id so it is globally unique across sources.
|
|||
|
|
id_offset = source_idx * TRACK_ID_OFFSET_PER_SOURCE
|
|||
|
|
for det in frame.get('detections', []):
|
|||
|
|
tagged = dict(det)
|
|||
|
|
tagged['lane_assignment'] = source_idx
|
|||
|
|
if 'track_id' in tagged and tagged['track_id'] is not None:
|
|||
|
|
tagged['track_id'] = tagged['track_id'] + id_offset
|
|||
|
|
merged_detections.append(tagged)
|
|||
|
|
|
|||
|
|
# Collect per-source tracking_stats
|
|||
|
|
if 'tracking_stats' in frame:
|
|||
|
|
merged_stats[source_name] = frame['tracking_stats']
|
|||
|
|
|
|||
|
|
merged_frame = {
|
|||
|
|
'image_name': image_name,
|
|||
|
|
'detections': merged_detections,
|
|||
|
|
}
|
|||
|
|
if merged_stats:
|
|||
|
|
merged_frame['tracking_stats'] = merged_stats
|
|||
|
|
|
|||
|
|
merged_frames.append(merged_frame)
|
|||
|
|
|
|||
|
|
return merged_frames
|
|||
|
|
|
|||
|
|
|
|||
|
|
def write_merged_output(merged, output_path):
|
|||
|
|
"""Persist merged frames to disk."""
|
|||
|
|
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(merged, f, indent=2, ensure_ascii=False)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def print_merge_summary(merged, source_map, output_path):
|
|||
|
|
"""Print aggregate detection counts per source for one merged output."""
|
|||
|
|
idx_to_name = {i: s for i, (s, _) in enumerate(source_map)}
|
|||
|
|
total_by_source = {s: 0 for s, _ in source_map}
|
|||
|
|
for frame in merged:
|
|||
|
|
for det in frame.get('detections', []):
|
|||
|
|
lane = det.get('lane_assignment')
|
|||
|
|
src = idx_to_name.get(lane)
|
|||
|
|
if src is not None:
|
|||
|
|
total_by_source[src] += 1
|
|||
|
|
|
|||
|
|
print(f"\nMerge complete: {len(merged)} frame(s) written to {output_path}")
|
|||
|
|
for i, (src, _) in enumerate(source_map):
|
|||
|
|
print(f" lane_assignment={i} ({src}): {total_by_source[src]} detection(s) total")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_case(roi0_path, roi1_path, merge_path, output_path, source_names, extra_sources=None, verbose=True):
|
|||
|
|
"""Merge one case from explicit input paths."""
|
|||
|
|
source_map = build_source_map(
|
|||
|
|
roi0_path,
|
|||
|
|
roi1_path,
|
|||
|
|
merge_path,
|
|||
|
|
source_names,
|
|||
|
|
extra_sources=extra_sources,
|
|||
|
|
verbose=verbose,
|
|||
|
|
)
|
|||
|
|
if not source_map:
|
|||
|
|
if verbose:
|
|||
|
|
print("Error: No valid tracking files were loaded. Nothing to merge.")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
if verbose:
|
|||
|
|
print(f"Merging {len(source_map)} source(s): {[s for s, _ in source_map]}")
|
|||
|
|
merged = merge_tracking_results(source_map)
|
|||
|
|
write_merged_output(merged, output_path)
|
|||
|
|
if verbose:
|
|||
|
|
print_merge_summary(merged, source_map, output_path)
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_case_dirs(results_root, model_name):
|
|||
|
|
"""Return all case directories that contain roi1.json under results_root/model_name."""
|
|||
|
|
search_root = Path(results_root) / model_name
|
|||
|
|
if not search_root.exists():
|
|||
|
|
print(f"Error: model directory does not exist: {search_root}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
return sorted(path.parent for path in search_root.rglob('roi1.json'))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_cases_in_batch(results_root, model_name, source_names, output_name):
|
|||
|
|
"""Scan case directories in Python and merge each one sequentially."""
|
|||
|
|
case_dirs = find_case_dirs(results_root, model_name)
|
|||
|
|
if not case_dirs:
|
|||
|
|
print("Error: No roi1.json files were found. Nothing to merge.")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
print(f"Found {len(case_dirs)} case(s) under {Path(results_root) / model_name}")
|
|||
|
|
merged_cases = 0
|
|||
|
|
failed_cases = 0
|
|||
|
|
|
|||
|
|
for case_dir in case_dirs:
|
|||
|
|
print("")
|
|||
|
|
print("==========================================")
|
|||
|
|
print(f"Case : {case_dir}")
|
|||
|
|
print(f"Output : {case_dir / output_name}")
|
|||
|
|
print("==========================================")
|
|||
|
|
|
|||
|
|
ok = merge_case(
|
|||
|
|
roi0_path=case_dir / 'roi0.json',
|
|||
|
|
roi1_path=case_dir / 'roi1.json',
|
|||
|
|
merge_path=case_dir / 'merge.json',
|
|||
|
|
output_path=case_dir / output_name,
|
|||
|
|
source_names=source_names,
|
|||
|
|
)
|
|||
|
|
if ok:
|
|||
|
|
merged_cases += 1
|
|||
|
|
else:
|
|||
|
|
failed_cases += 1
|
|||
|
|
|
|||
|
|
print("")
|
|||
|
|
print("Batch merge complete")
|
|||
|
|
print(f" merged_cases: {merged_cases}")
|
|||
|
|
print(f" failed_cases: {failed_cases}")
|
|||
|
|
return failed_cases == 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_extra_source_specs(specs):
|
|||
|
|
"""Parse CLI specs of the form NAME=PATH for additional tracking sources."""
|
|||
|
|
parsed = []
|
|||
|
|
for spec in specs or []:
|
|||
|
|
if '=' not in spec:
|
|||
|
|
raise ValueError(f"Invalid --extra-source {spec!r}; expected NAME=PATH")
|
|||
|
|
source_name, source_path = spec.split('=', 1)
|
|||
|
|
source_name = source_name.strip()
|
|||
|
|
source_path = source_path.strip()
|
|||
|
|
if not source_name or not source_path:
|
|||
|
|
raise ValueError(f"Invalid --extra-source {spec!r}; expected NAME=PATH")
|
|||
|
|
parsed.append((source_name, source_path))
|
|||
|
|
return parsed
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(
|
|||
|
|
description='Merge per-source tracking JSONs into a single combined JSON.'
|
|||
|
|
)
|
|||
|
|
parser.add_argument('--roi0', type=str, default=None,
|
|||
|
|
help='Path to ROI0 tracking JSON (output of track_objects.py)')
|
|||
|
|
parser.add_argument('--roi1', type=str, default=None,
|
|||
|
|
help='Path to ROI1 tracking JSON (output of track_objects.py)')
|
|||
|
|
parser.add_argument('--merge', type=str, default=None,
|
|||
|
|
help='Path to merged-ROI tracking JSON (output of track_objects.py)')
|
|||
|
|
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('--output-name', type=str, default='combined_tracking.json',
|
|||
|
|
help='Batch mode output filename written inside each case directory')
|
|||
|
|
parser.add_argument('--source-names', type=str, nargs=3,
|
|||
|
|
default=['roi0', 'roi1', 'merge'],
|
|||
|
|
metavar=('ROI0_NAME', 'ROI1_NAME', 'MERGE_NAME'),
|
|||
|
|
help='Custom labels for the three sources (default: roi0 roi1 merge)')
|
|||
|
|
parser.add_argument('--extra-source', type=str, action='append', default=[],
|
|||
|
|
metavar='NAME=PATH',
|
|||
|
|
help='Additional tracking source to append in single-case mode. Can be passed multiple times.')
|
|||
|
|
parser.add_argument('--output', type=str, default=None,
|
|||
|
|
help='Output path for the combined tracking JSON in single-case mode')
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
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')
|
|||
|
|
merge_cases_in_batch(
|
|||
|
|
results_root=args.results_root,
|
|||
|
|
model_name=args.model_name,
|
|||
|
|
source_names=args.source_names,
|
|||
|
|
output_name=args.output_name,
|
|||
|
|
)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
if args.output is None:
|
|||
|
|
parser.error('--output is required in single-case mode')
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
extra_sources = parse_extra_source_specs(args.extra_source)
|
|||
|
|
except ValueError as exc:
|
|||
|
|
parser.error(str(exc))
|
|||
|
|
|
|||
|
|
merge_case(
|
|||
|
|
roi0_path=args.roi0,
|
|||
|
|
roi1_path=args.roi1,
|
|||
|
|
merge_path=args.merge,
|
|||
|
|
output_path=args.output,
|
|||
|
|
source_names=args.source_names,
|
|||
|
|
extra_sources=extra_sources,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|