单目3D初始代码

tools/temporal_analysis/split_tracking_json.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""Split tracking JSON into one JSON file per frame while preserving track IDs.
+
+Input format:
+    [
+      {
+        "image_name": "1741824631_0",
+        "detections": [
+          {
+            "class_id": 0,
+            "confidence": 0.91,
+            "bbox": [x1, y1, x2, y2],
+            "track_id": 7,
+            ...
+          }
+        ]
+      },
+      ...
+    ]
+
+Output format (one file per frame, evaluator-friendly):
+    {
+      "0": {
+        "type": "0",
+        "type_name": "vehicle",
+        "score": "0.91",
+        "box2d": [x1, y1, x2, y2],
+        "xyzlhwyaw": [...],
+        "face_cls": "front",
+        "cut_cls": "0",
+        "cut_cls_name": "nocut",
+        "track_id": 7
+      },
+      "1": { ... }
+    }
+
+Usage:
+    python tools/temporal_analysis/split_tracking_json.py \
+        --input /path/to/roi0.json \
+        --output-dir /path/to/json_results
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+
+from merge_tracking_results import normalize_image_name
+
+
+CLASS_ID_TO_NAME = {
+    0: "vehicle",
+    1: "pedestrian",
+    2: "bicycle",
+    3: "rider",
+    4: "roadblock",
+    5: "head",
+    6: "tsr",
+    7: "guideboard",
+    8: "plate",
+    9: "wheel",
+    10: "tl_border",
+    11: "tl_wick",
+    12: "tl_num",
+    13: "tricycle",
+}
+
+
+def format_scalar(value):
+    """Keep ints as ints, convert everything else to plain JSON numbers/strings."""
+    if value is None:
+        return value
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        return float(value)
+    return value
+
+
+def normalize_detection(det, keep_all_fields=False):
+    """Convert one tracking detection to evaluator-friendly single-frame JSON schema."""
+    class_id = int(det.get("class_id", -1))
+    bbox = det.get("bbox", [])
+    confidence = det.get("confidence", 0.0)
+
+    item = {
+        "type": str(class_id),
+        "type_name": det.get("type_name") or CLASS_ID_TO_NAME.get(class_id, ""),
+        "score": confidence,
+        "box2d": [format_scalar(v) for v in bbox[:4]],
+        "face_cls": det.get("face_cls", "none"),
+        "cut_cls": str(det.get("cut_cls", -1)),
+        "cut_cls_name": det.get("cut_cls_name", "none"),
+    }
+
+    if "track_id" in det and det.get("track_id") is not None:
+        item["track_id"] = int(det["track_id"])
+    if "lane_assignment" in det and det.get("lane_assignment") is not None:
+        item["lane_assignment"] = int(det["lane_assignment"])
+    if "frameId" in det and det.get("frameId") is not None:
+        item["frameId"] = det["frameId"]
+    if "timestamp" in det and det.get("timestamp") is not None:
+        item["timestamp"] = format_scalar(det["timestamp"])
+    if "version" in det and det.get("version") is not None:
+        item["version"] = det["version"]
+    if "anchor" in det and det.get("anchor") is not None:
+        item["anchor"] = det["anchor"]
+
+    if "roi_id" in det and det.get("roi_id") is not None:
+        item["roi_id"] = int(det["roi_id"])
+
+    object_3d = det.get("object_3d")
+    if object_3d is not None:
+        item["xyzlhwyaw"] = [format_scalar(v) for v in object_3d[:7]]
+
+    object_3d_ego = det.get("object_3d_ego")
+    if object_3d_ego is not None:
+        item["xyzlhwyaw_ego"] = [format_scalar(v) for v in object_3d_ego[:7]]
+
+    if keep_all_fields:
+        for key, value in det.items():
+            if key in item or key in {"class_id", "confidence", "bbox", "object_3d", "object_3d_ego"}:
+                continue
+            item[key] = value
+
+    return item
+
+
+def frame_output_name(frame):
+    """Choose output filename stem for one frame."""
+    image_name = frame.get("image_name") or ""
+    stem = Path(image_name).stem if image_name else ""
+    return normalize_image_name(stem) if stem else "unknown"
+
+
+def build_output_data(frame, keep_all_fields=False):
+    """Convert one tracking frame to per-frame output payload."""
+    output_data = {}
+    for idx, det in enumerate(frame.get("detections", [])):
+        output_data[str(idx)] = normalize_detection(det, keep_all_fields=keep_all_fields)
+    return output_data
+
+
+def write_one_frame(frame, output_dir, keep_all_fields=False, pretty=False):
+    """Write one output JSON file and return its path."""
+    frame_name = frame_output_name(frame)
+    output_data = build_output_data(frame, keep_all_fields=keep_all_fields)
+    output_path = output_dir / f"{frame_name}.json"
+
+    dump_kwargs = {
+        "ensure_ascii": False,
+    }
+    if pretty:
+        dump_kwargs["indent"] = 2
+    else:
+        dump_kwargs["separators"] = (",", ":")
+
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(output_data, f, **dump_kwargs)
+    return output_path
+
+
+def split_tracking_json(input_path, output_dir, keep_all_fields=False, num_workers=1, pretty=False):
+    """Split a tracking JSON list into one frame file per item."""
+    input_path = Path(input_path)
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(input_path, "r", encoding="utf-8") as f:
+        frames = json.load(f)
+
+    if not isinstance(frames, list):
+        raise ValueError(f"Expected list of frames in {input_path}, got {type(frames).__name__}")
+
+    max_workers = max(1, int(num_workers))
+    if max_workers == 1 or len(frames) <= 1:
+        for frame in frames:
+            write_one_frame(frame, output_dir, keep_all_fields=keep_all_fields, pretty=pretty)
+        return len(frames)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(
+                write_one_frame,
+                frame,
+                output_dir,
+                keep_all_fields,
+                pretty,
+            )
+            for frame in frames
+        ]
+        for future in futures:
+            future.result()
+
+    return len(frames)
+
+
+def default_output_dir(input_path):
+    """Return a default output directory next to the input JSON."""
+    input_path = Path(input_path)
+    return input_path.parent / f"{input_path.stem}_split"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Split tracking JSON into evaluator-friendly per-frame JSON files."
+    )
+    parser.add_argument("--input", required=True, help="Path to roi0.json / tracking.json / merge.json")
+    parser.add_argument("--output-dir", default=None,
+                        help="Directory to write per-frame JSON files (default: <input_stem>_split)")
+    parser.add_argument("--keep-all-fields", action="store_true",
+                        help="Keep extra tracking fields beyond the evaluator-friendly schema")
+    parser.add_argument("--num-workers", type=int, default=max(1, min(16, os.cpu_count() or 1)),
+                        help="Number of parallel workers for per-frame JSON writing (default: %(default)s)")
+    parser.add_argument("--pretty", action="store_true",
+                        help="Write indented JSON for readability; default uses compact JSON for speed")
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir) if args.output_dir else default_output_dir(args.input)
+    written = split_tracking_json(
+        input_path=args.input,
+        output_dir=output_dir,
+        keep_all_fields=args.keep_all_fields,
+        num_workers=args.num_workers,
+        pretty=args.pretty,
+    )
+
+    print(f"Input tracking JSON : {args.input}")
+    print(f"Output directory    : {output_dir}")
+    print(f"Frames written      : {written}")
+    print(f"Worker count        : {args.num_workers}")
+    print(f"Pretty JSON         : {args.pretty}")
+
+
+if __name__ == "__main__":
+    main()