686 lines
24 KiB
Python
686 lines
24 KiB
Python
|
"""
|
|||
|
|
Visualize tracking results by drawing 2D boxes on images.
|
|||
|
|
|
|||
|
|
Usage:
|
|||
|
|
# 从图像文件夹可视化
|
|||
|
|
python visualize_tracking_boxes.py \
|
|||
|
|
--tracking runs/val_viz/exp/tracking.json \
|
|||
|
|
--images path/to/images \
|
|||
|
|
--output runs/val_viz/exp/tracked_images
|
|||
|
|
|
|||
|
|
# 生成视频
|
|||
|
|
python visualize_tracking_boxes.py \
|
|||
|
|
--tracking runs/val_viz/exp/tracking.json \
|
|||
|
|
--images path/to/images \
|
|||
|
|
--output runs/val_viz/exp/tracked.mp4 \
|
|||
|
|
--save-video
|
|||
|
|
|
|||
|
|
# 只可视化前100帧
|
|||
|
|
python visualize_tracking_boxes.py \
|
|||
|
|
--tracking tracking.json \
|
|||
|
|
--images images/ \
|
|||
|
|
--output output/ \
|
|||
|
|
--max-frames 100
|
|||
|
|
|
|||
|
|
# 按 frameId 对齐原始 camera4.bin，并导出单帧图片
|
|||
|
|
python visualize_tracking_boxes.py \
|
|||
|
|
--tracking merge.json \
|
|||
|
|
--images /path/to/case/sigmastar.1/camera4.bin \
|
|||
|
|
--output tracking_vis_raw \
|
|||
|
|
--align-by-frame-id
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import json
|
|||
|
|
from pathlib import Path
|
|||
|
|
import re
|
|||
|
|
import sys
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
import colorsys
|
|||
|
|
|
|||
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|||
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|||
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|||
|
|
|
|||
|
|
from tools.model_inference.adapters.video_dir_inference_utils import (
|
|||
|
|
iter_video_case_frames,
|
|||
|
|
read_video_frame_index,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_colors(n):
|
|||
|
|
"""生成n个区分度高的颜色"""
|
|||
|
|
colors = []
|
|||
|
|
for i in range(n):
|
|||
|
|
hue = i / n
|
|||
|
|
saturation = 0.9
|
|||
|
|
value = 0.9
|
|||
|
|
rgb = colorsys.hsv_to_rgb(hue, saturation, value)
|
|||
|
|
colors.append(tuple(int(x * 255) for x in rgb))
|
|||
|
|
return colors
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_color_for_track(track_id, color_map):
|
|||
|
|
"""为track_id分配一个持久的颜色"""
|
|||
|
|
if track_id not in color_map:
|
|||
|
|
# 使用track_id作为种子生成颜色
|
|||
|
|
np.random.seed(track_id)
|
|||
|
|
color = (
|
|||
|
|
np.random.randint(50, 255),
|
|||
|
|
np.random.randint(50, 255),
|
|||
|
|
np.random.randint(50, 255)
|
|||
|
|
)
|
|||
|
|
color_map[track_id] = color
|
|||
|
|
return color_map[track_id]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def should_visualize_detection(det, class_id_filter=None, track_ids_filter=None):
|
|||
|
|
"""判断检测是否应该被可视化
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
det: 检测结果字典
|
|||
|
|
class_id_filter: 指定的类别ID，None表示所有类别
|
|||
|
|
track_ids_filter: 指定的track ID列表，None表示所有ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 是否应该可视化该检测
|
|||
|
|
"""
|
|||
|
|
# 检查class_id过滤
|
|||
|
|
if class_id_filter is not None and det.get('class_id') != class_id_filter:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 检查track_id过滤
|
|||
|
|
if track_ids_filter is not None:
|
|||
|
|
track_id = det.get('track_id')
|
|||
|
|
if track_id not in track_ids_filter:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def safe_int(value):
|
|||
|
|
"""Best-effort integer conversion."""
|
|||
|
|
if value is None:
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
return int(value)
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_frame_id_from_image_name(image_name):
|
|||
|
|
"""Best-effort frame-id extraction from tracking image_name."""
|
|||
|
|
stem = Path(str(image_name or "")).stem
|
|||
|
|
match = re.search(r"_(\d+)_(\d+)(?:_merged)?$", stem)
|
|||
|
|
if match is not None:
|
|||
|
|
return int(match.group(1))
|
|||
|
|
|
|||
|
|
match = re.search(r"_(\d+)(?:_merged)?$", stem)
|
|||
|
|
if match is not None:
|
|||
|
|
return int(match.group(1))
|
|||
|
|
|
|||
|
|
match = re.search(r"(\d+)", stem)
|
|||
|
|
if match is not None:
|
|||
|
|
return int(match.group(1))
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_tracking_frame_id(frame_data, fallback_idx=None):
|
|||
|
|
"""Resolve one frame's frameId from detections, frame_info, or image_name."""
|
|||
|
|
for det in frame_data.get("detections", []):
|
|||
|
|
for key in ("frameId", "frame_id"):
|
|||
|
|
frame_id = safe_int(det.get(key))
|
|||
|
|
if frame_id is not None:
|
|||
|
|
return frame_id
|
|||
|
|
|
|||
|
|
frame_info = frame_data.get("frame_info")
|
|||
|
|
if isinstance(frame_info, dict):
|
|||
|
|
for key in ("original_frame_id", "frame_id", "frameId"):
|
|||
|
|
frame_id = safe_int(frame_info.get(key))
|
|||
|
|
if frame_id is not None:
|
|||
|
|
return frame_id
|
|||
|
|
|
|||
|
|
frame_id = extract_frame_id_from_image_name(frame_data.get("image_name"))
|
|||
|
|
if frame_id is not None:
|
|||
|
|
return frame_id
|
|||
|
|
return fallback_idx
|
|||
|
|
|
|||
|
|
|
|||
|
|
def resolve_detection_class_name(det, class_id, class_names):
|
|||
|
|
"""Prefer detection-provided type_name before falling back to static names."""
|
|||
|
|
type_name = str(det.get("type_name") or "").strip()
|
|||
|
|
if type_name:
|
|||
|
|
return type_name
|
|||
|
|
return class_names.get(class_id, f"Class{class_id}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_output_frame_name(frame_data, fallback_idx):
|
|||
|
|
"""Build a stable output filename for one visualized frame."""
|
|||
|
|
frame_name = frame_data.get("image_name", f"frame_{fallback_idx:06d}")
|
|||
|
|
stem = Path(str(frame_name)).stem
|
|||
|
|
return f"{stem}.jpg"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def draw_box_with_label(img, bbox, track_id, class_id, confidence, color, class_names, class_name_override=None):
|
|||
|
|
"""在图像上绘制边界框和标签"""
|
|||
|
|
x1, y1, x2, y2 = map(int, bbox)
|
|||
|
|
|
|||
|
|
# 绘制边界框
|
|||
|
|
thickness = 2
|
|||
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness)
|
|||
|
|
|
|||
|
|
# 准备标签文本
|
|||
|
|
class_name = class_name_override or class_names.get(class_id, f'Class{class_id}')
|
|||
|
|
label = f'ID:{track_id} {class_name} {confidence:.2f}'
|
|||
|
|
|
|||
|
|
# 计算标签背景大小
|
|||
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|||
|
|
font_scale = 0.6
|
|||
|
|
font_thickness = 2
|
|||
|
|
(text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, font_thickness)
|
|||
|
|
|
|||
|
|
# 绘制标签背景
|
|||
|
|
label_y1 = max(y1 - text_height - 10, 0)
|
|||
|
|
label_y2 = y1
|
|||
|
|
cv2.rectangle(img, (x1, label_y1), (x1 + text_width + 10, label_y2), color, -1)
|
|||
|
|
|
|||
|
|
# 绘制标签文本
|
|||
|
|
text_x = x1 + 5
|
|||
|
|
text_y = label_y1 + text_height + 5
|
|||
|
|
cv2.putText(img, label, (text_x, text_y), font, font_scale, (255, 255, 255), font_thickness)
|
|||
|
|
|
|||
|
|
return img
|
|||
|
|
|
|||
|
|
|
|||
|
|
def render_tracking_frame(
|
|||
|
|
img,
|
|||
|
|
frame_data,
|
|||
|
|
frame_idx,
|
|||
|
|
total_frames,
|
|||
|
|
color_map,
|
|||
|
|
class_names,
|
|||
|
|
trajectory_history=None,
|
|||
|
|
class_id_filter=None,
|
|||
|
|
track_ids_filter=None,
|
|||
|
|
):
|
|||
|
|
"""Render one tracking frame on top of an image."""
|
|||
|
|
detections = frame_data.get('detections', [])
|
|||
|
|
visualized_count = 0
|
|||
|
|
for det in detections:
|
|||
|
|
track_id = det.get('track_id')
|
|||
|
|
if track_id is None:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if not should_visualize_detection(det, class_id_filter, track_ids_filter):
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
visualized_count += 1
|
|||
|
|
bbox = det['bbox']
|
|||
|
|
class_id = det['class_id']
|
|||
|
|
confidence = det['confidence']
|
|||
|
|
color = get_color_for_track(track_id, color_map)
|
|||
|
|
class_name = resolve_detection_class_name(det, class_id, class_names)
|
|||
|
|
|
|||
|
|
img = draw_box_with_label(
|
|||
|
|
img,
|
|||
|
|
bbox,
|
|||
|
|
track_id,
|
|||
|
|
class_id,
|
|||
|
|
confidence,
|
|||
|
|
color,
|
|||
|
|
class_names,
|
|||
|
|
class_name_override=class_name,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if trajectory_history is not None:
|
|||
|
|
center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)
|
|||
|
|
if track_id not in trajectory_history:
|
|||
|
|
trajectory_history[track_id] = []
|
|||
|
|
trajectory_history[track_id].append(center)
|
|||
|
|
|
|||
|
|
if len(trajectory_history[track_id]) > 1:
|
|||
|
|
points = np.array(trajectory_history[track_id], dtype=np.int32)
|
|||
|
|
for i in range(len(points) - 1):
|
|||
|
|
cv2.line(img, tuple(points[i]), tuple(points[i + 1]), color, 2)
|
|||
|
|
|
|||
|
|
if class_id_filter is not None or track_ids_filter is not None:
|
|||
|
|
info_text = f"Frame: {frame_idx + 1}/{total_frames} | Showing: {visualized_count}/{len(detections)}"
|
|||
|
|
else:
|
|||
|
|
info_text = f"Frame: {frame_idx + 1}/{total_frames} | Tracks: {len(detections)}"
|
|||
|
|
cv2.putText(img, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
|
|||
|
|
0.8, (0, 255, 0), 2)
|
|||
|
|
return img, visualized_count
|
|||
|
|
|
|||
|
|
|
|||
|
|
def visualize_tracking_on_raw_video_by_frame_id(
|
|||
|
|
tracking_data,
|
|||
|
|
video_path,
|
|||
|
|
output_path,
|
|||
|
|
max_frames=None,
|
|||
|
|
show_trajectory=False,
|
|||
|
|
class_id_filter=None,
|
|||
|
|
track_ids_filter=None,
|
|||
|
|
):
|
|||
|
|
"""Visualize tracking results on raw camera4.bin frames using frameId alignment."""
|
|||
|
|
video_path = Path(video_path)
|
|||
|
|
output_path = Path(output_path)
|
|||
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
class_names = {
|
|||
|
|
0: 'Car',
|
|||
|
|
1: 'Pedestrian',
|
|||
|
|
2: 'Tricycle',
|
|||
|
|
3: 'Cyclist'
|
|||
|
|
}
|
|||
|
|
color_map = {}
|
|||
|
|
trajectory_history = {} if show_trajectory else None
|
|||
|
|
|
|||
|
|
selected_tracking_data = tracking_data[:max_frames] if max_frames else tracking_data
|
|||
|
|
frame_id_to_tracking = {}
|
|||
|
|
ordered_frame_ids = []
|
|||
|
|
skipped_duplicates = []
|
|||
|
|
unresolved_frames = []
|
|||
|
|
for idx, frame_data in enumerate(selected_tracking_data):
|
|||
|
|
frame_id = extract_tracking_frame_id(frame_data, fallback_idx=idx)
|
|||
|
|
if frame_id is None:
|
|||
|
|
unresolved_frames.append(frame_data.get("image_name", f"frame_{idx:06d}"))
|
|||
|
|
continue
|
|||
|
|
if frame_id in frame_id_to_tracking:
|
|||
|
|
skipped_duplicates.append(frame_id)
|
|||
|
|
continue
|
|||
|
|
frame_id_to_tracking[frame_id] = (idx, frame_data)
|
|||
|
|
ordered_frame_ids.append(frame_id)
|
|||
|
|
|
|||
|
|
if not frame_id_to_tracking:
|
|||
|
|
print("Error: no frameIds could be resolved from tracking data")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
frame_index_payload = read_video_frame_index(video_path)
|
|||
|
|
total_targets = len(ordered_frame_ids)
|
|||
|
|
matched_frame_ids = set()
|
|||
|
|
|
|||
|
|
print(f"Reading raw video with frameId alignment: {video_path}")
|
|||
|
|
print(f"Saving aligned images to: {output_path}")
|
|||
|
|
if unresolved_frames:
|
|||
|
|
print(f"Warning: skipped {len(unresolved_frames)} frame(s) without resolvable frameId")
|
|||
|
|
if skipped_duplicates:
|
|||
|
|
print(f"Warning: skipped duplicate frameId(s): {sorted(set(skipped_duplicates))[:10]}")
|
|||
|
|
|
|||
|
|
progress = tqdm(total=total_targets, desc="Visualizing")
|
|||
|
|
for read_frame_index, frame_bgr, _, frame_info in iter_video_case_frames(
|
|||
|
|
video_path,
|
|||
|
|
frame_index_payload=frame_index_payload,
|
|||
|
|
frame_stride=1,
|
|||
|
|
max_frames=0,
|
|||
|
|
):
|
|||
|
|
aligned_frame_id = safe_int(None if frame_info is None else frame_info.get("frame_id"))
|
|||
|
|
if aligned_frame_id is None:
|
|||
|
|
aligned_frame_id = read_frame_index
|
|||
|
|
|
|||
|
|
if aligned_frame_id not in frame_id_to_tracking:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
tracking_idx, frame_data = frame_id_to_tracking[aligned_frame_id]
|
|||
|
|
rendered, _ = render_tracking_frame(
|
|||
|
|
frame_bgr.copy(),
|
|||
|
|
frame_data,
|
|||
|
|
tracking_idx,
|
|||
|
|
total_targets,
|
|||
|
|
color_map,
|
|||
|
|
class_names,
|
|||
|
|
trajectory_history=trajectory_history,
|
|||
|
|
class_id_filter=class_id_filter,
|
|||
|
|
track_ids_filter=track_ids_filter,
|
|||
|
|
)
|
|||
|
|
output_file = output_path / build_output_frame_name(frame_data, tracking_idx)
|
|||
|
|
cv2.imwrite(str(output_file), rendered)
|
|||
|
|
matched_frame_ids.add(aligned_frame_id)
|
|||
|
|
progress.update(1)
|
|||
|
|
|
|||
|
|
if len(matched_frame_ids) >= total_targets:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
progress.close()
|
|||
|
|
|
|||
|
|
missing_frame_ids = [frame_id for frame_id in ordered_frame_ids if frame_id not in matched_frame_ids]
|
|||
|
|
print(f"\n✓ Visualization complete!")
|
|||
|
|
print(f" Images saved to: {output_path}")
|
|||
|
|
print(f" Matched frames: {len(matched_frame_ids)}/{total_targets}")
|
|||
|
|
print(f" Total tracks visualized: {len(color_map)}")
|
|||
|
|
if missing_frame_ids:
|
|||
|
|
preview = ", ".join(str(frame_id) for frame_id in missing_frame_ids[:10])
|
|||
|
|
suffix = "..." if len(missing_frame_ids) > 10 else ""
|
|||
|
|
print(f" Missing frameIds: {preview}{suffix}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def visualize_tracking_on_images(tracking_data, image_source, output_path,
|
|||
|
|
save_video=False, fps=25, max_frames=None,
|
|||
|
|
show_trajectory=False, class_id_filter=None,
|
|||
|
|
track_ids_filter=None, align_by_frame_id=False):
|
|||
|
|
"""将跟踪结果可视化到图像上
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
tracking_data: 跟踪结果数据（从tracking.json加载）
|
|||
|
|
image_source: 图像来源（文件夹路径或视频文件路径）
|
|||
|
|
output_path: 输出路径（文件夹或视频文件）
|
|||
|
|
save_video: 是否保存为视频
|
|||
|
|
fps: 视频帧率
|
|||
|
|
max_frames: 最大处理帧数
|
|||
|
|
show_trajectory: 是否显示轨迹线
|
|||
|
|
class_id_filter: 指定类别ID，None表示所有类别
|
|||
|
|
track_ids_filter: 指定track ID集合，None表示所有ID
|
|||
|
|
align_by_frame_id: 是否按 frameId 对齐原始视频帧
|
|||
|
|
"""
|
|||
|
|
image_source = Path(image_source)
|
|||
|
|
output_path = Path(output_path)
|
|||
|
|
|
|||
|
|
if align_by_frame_id:
|
|||
|
|
if save_video:
|
|||
|
|
print("Error: --align-by-frame-id only supports image-directory output")
|
|||
|
|
return
|
|||
|
|
if not image_source.is_file():
|
|||
|
|
print("Error: --align-by-frame-id requires a video file such as camera4.bin")
|
|||
|
|
return
|
|||
|
|
visualize_tracking_on_raw_video_by_frame_id(
|
|||
|
|
tracking_data,
|
|||
|
|
image_source,
|
|||
|
|
output_path,
|
|||
|
|
max_frames=max_frames,
|
|||
|
|
show_trajectory=show_trajectory,
|
|||
|
|
class_id_filter=class_id_filter,
|
|||
|
|
track_ids_filter=track_ids_filter,
|
|||
|
|
)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 类别名称
|
|||
|
|
class_names = {
|
|||
|
|
0: 'Car',
|
|||
|
|
1: 'Pedestrian',
|
|||
|
|
2: 'Tricycle',
|
|||
|
|
3: 'Cyclist'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 颜色映射
|
|||
|
|
color_map = {}
|
|||
|
|
|
|||
|
|
# 轨迹历史（用于绘制轨迹线）
|
|||
|
|
trajectory_history = {} if show_trajectory else None
|
|||
|
|
|
|||
|
|
# 检查图像源类型
|
|||
|
|
VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.bin'}
|
|||
|
|
if image_source.is_file() and image_source.suffix.lower() in VIDEO_EXTENSIONS:
|
|||
|
|
# 从视频读取
|
|||
|
|
cap = cv2.VideoCapture(str(image_source))
|
|||
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|||
|
|
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|||
|
|
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|||
|
|
print(f"Reading from video: {image_source}")
|
|||
|
|
print(f"Video info: {frame_width}x{frame_height}, {total_frames} frames")
|
|||
|
|
elif image_source.is_dir():
|
|||
|
|
# 从图像文件夹读取
|
|||
|
|
cap = None
|
|||
|
|
image_files = sorted(list(image_source.glob('*.jpg')) +
|
|||
|
|
list(image_source.glob('*.png')) +
|
|||
|
|
list(image_source.glob('*.jpeg')))
|
|||
|
|
if not image_files:
|
|||
|
|
print(f"Error: No images found in {image_source}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 读取第一张图像获取尺寸
|
|||
|
|
first_img = cv2.imread(str(image_files[0]))
|
|||
|
|
if first_img is None:
|
|||
|
|
print(f"Error: Cannot read image {image_files[0]}")
|
|||
|
|
return
|
|||
|
|
frame_height, frame_width = first_img.shape[:2]
|
|||
|
|
total_frames = len(image_files)
|
|||
|
|
print(f"Reading from image folder: {image_source}")
|
|||
|
|
print(f"Found {total_frames} images, size: {frame_width}x{frame_height}")
|
|||
|
|
else:
|
|||
|
|
print(f"Error: Invalid image source: {image_source}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 限制处理帧数
|
|||
|
|
if max_frames:
|
|||
|
|
total_frames = min(total_frames, max_frames)
|
|||
|
|
tracking_data = tracking_data[:total_frames]
|
|||
|
|
|
|||
|
|
# 准备输出
|
|||
|
|
if save_video:
|
|||
|
|
if not str(output_path).endswith(('.mp4', '.avi')):
|
|||
|
|
output_path = output_path.with_suffix('.mp4')
|
|||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|||
|
|
video_writer = cv2.VideoWriter(str(output_path), fourcc, fps, (frame_width, frame_height))
|
|||
|
|
print(f"Saving video to: {output_path}")
|
|||
|
|
else:
|
|||
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
video_writer = None
|
|||
|
|
print(f"Saving images to: {output_path}")
|
|||
|
|
|
|||
|
|
# 处理每一帧
|
|||
|
|
print(f"\nProcessing {len(tracking_data)} frames...")
|
|||
|
|
for frame_idx, frame_data in enumerate(tqdm(tracking_data, desc="Visualizing")):
|
|||
|
|
# 读取图像
|
|||
|
|
if cap is not None:
|
|||
|
|
ret, img = cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
print(f"Warning: Cannot read frame {frame_idx}")
|
|||
|
|
break
|
|||
|
|
else:
|
|||
|
|
if frame_idx >= len(image_files):
|
|||
|
|
break
|
|||
|
|
img = cv2.imread(str(image_files[frame_idx]))
|
|||
|
|
if img is None:
|
|||
|
|
print(f"Warning: Cannot read image {image_files[frame_idx]}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
img, _ = render_tracking_frame(
|
|||
|
|
img,
|
|||
|
|
frame_data,
|
|||
|
|
frame_idx,
|
|||
|
|
len(tracking_data),
|
|||
|
|
color_map,
|
|||
|
|
class_names,
|
|||
|
|
trajectory_history=trajectory_history,
|
|||
|
|
class_id_filter=class_id_filter,
|
|||
|
|
track_ids_filter=track_ids_filter,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 保存或写入视频
|
|||
|
|
if video_writer:
|
|||
|
|
video_writer.write(img)
|
|||
|
|
else:
|
|||
|
|
frame_name = frame_data.get('image_name', f'frame_{frame_idx:06d}.jpg')
|
|||
|
|
if not frame_name.endswith(('.jpg', '.png', '.jpeg')):
|
|||
|
|
frame_name = f'{Path(frame_name).stem}.jpg'
|
|||
|
|
output_file = output_path / frame_name.replace('.png', '.jpg')
|
|||
|
|
cv2.imwrite(str(output_file), img)
|
|||
|
|
|
|||
|
|
# 清理
|
|||
|
|
if cap is not None:
|
|||
|
|
cap.release()
|
|||
|
|
if video_writer is not None:
|
|||
|
|
video_writer.release()
|
|||
|
|
|
|||
|
|
print(f"\n✓ Visualization complete!")
|
|||
|
|
if save_video:
|
|||
|
|
print(f" Video saved to: {output_path}")
|
|||
|
|
else:
|
|||
|
|
print(f" Images saved to: {output_path}")
|
|||
|
|
print(f" Total tracks visualized: {len(color_map)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def create_side_by_side_comparison(tracking_file1, tracking_file2, image_source,
|
|||
|
|
output_path, labels=None, max_frames=100):
|
|||
|
|
"""创建并排对比可视化
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
tracking_file1, tracking_file2: 两个跟踪结果文件
|
|||
|
|
image_source: 图像源
|
|||
|
|
output_path: 输出视频路径
|
|||
|
|
labels: 两个方法的标签
|
|||
|
|
max_frames: 最大处理帧数
|
|||
|
|
"""
|
|||
|
|
if labels is None:
|
|||
|
|
labels = ['Method 1', 'Method 2']
|
|||
|
|
|
|||
|
|
# 加载跟踪数据
|
|||
|
|
with open(tracking_file1, 'r') as f:
|
|||
|
|
data1 = json.load(f)
|
|||
|
|
with open(tracking_file2, 'r') as f:
|
|||
|
|
data2 = json.load(f)
|
|||
|
|
|
|||
|
|
image_source = Path(image_source)
|
|||
|
|
output_path = Path(output_path)
|
|||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
class_names = {0: 'Car', 1: 'Pedestrian', 2: 'Tricycle', 3: 'Cyclist'}
|
|||
|
|
color_map1, color_map2 = {}, {}
|
|||
|
|
|
|||
|
|
# 读取图像
|
|||
|
|
if image_source.is_dir():
|
|||
|
|
image_files = sorted(list(image_source.glob('*.jpg')) +
|
|||
|
|
list(image_source.glob('*.png')))
|
|||
|
|
first_img = cv2.imread(str(image_files[0]))
|
|||
|
|
h, w = first_img.shape[:2]
|
|||
|
|
else:
|
|||
|
|
cap = cv2.VideoCapture(str(image_source))
|
|||
|
|
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|||
|
|
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|||
|
|
cap.release()
|
|||
|
|
|
|||
|
|
# 创建视频写入器（并排宽度加倍）
|
|||
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|||
|
|
video_writer = cv2.VideoWriter(str(output_path), fourcc, 25, (w * 2, h))
|
|||
|
|
|
|||
|
|
print(f"Creating side-by-side comparison video...")
|
|||
|
|
frames_to_process = min(len(data1), len(data2), max_frames)
|
|||
|
|
|
|||
|
|
for frame_idx in tqdm(range(frames_to_process), desc="Processing"):
|
|||
|
|
# 读取原始图像
|
|||
|
|
if image_source.is_dir():
|
|||
|
|
img = cv2.imread(str(image_files[frame_idx]))
|
|||
|
|
else:
|
|||
|
|
cap = cv2.VideoCapture(str(image_source))
|
|||
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
|
|||
|
|
ret, img = cap.read()
|
|||
|
|
cap.release()
|
|||
|
|
if not ret:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
img1 = img.copy()
|
|||
|
|
img2 = img.copy()
|
|||
|
|
|
|||
|
|
# 绘制第一个跟踪结果
|
|||
|
|
for det in data1[frame_idx].get('detections', []):
|
|||
|
|
if 'track_id' in det:
|
|||
|
|
color = get_color_for_track(det['track_id'], color_map1)
|
|||
|
|
img1 = draw_box_with_label(img1, det['bbox'], det['track_id'],
|
|||
|
|
det['class_id'], det['confidence'], color, class_names)
|
|||
|
|
|
|||
|
|
# 绘制第二个跟踪结果
|
|||
|
|
for det in data2[frame_idx].get('detections', []):
|
|||
|
|
if 'track_id' in det:
|
|||
|
|
color = get_color_for_track(det['track_id'], color_map2)
|
|||
|
|
img2 = draw_box_with_label(img2, det['bbox'], det['track_id'],
|
|||
|
|
det['class_id'], det['confidence'], color, class_names)
|
|||
|
|
|
|||
|
|
# 添加标签
|
|||
|
|
cv2.putText(img1, labels[0], (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
|||
|
|
cv2.putText(img2, labels[1], (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
|||
|
|
|
|||
|
|
# 并排拼接
|
|||
|
|
combined = np.hstack([img1, img2])
|
|||
|
|
video_writer.write(combined)
|
|||
|
|
|
|||
|
|
video_writer.release()
|
|||
|
|
print(f"✓ Comparison video saved to: {output_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='Visualize tracking results on images')
|
|||
|
|
parser.add_argument('--tracking', type=str, required=True, help='Path to tracking.json file')
|
|||
|
|
parser.add_argument('--images', type=str, required=True, help='Path to images folder or video file')
|
|||
|
|
parser.add_argument('--output', type=str, required=True, help='Output path (folder or video file)')
|
|||
|
|
parser.add_argument('--save-video', action='store_true', help='Save as video instead of images')
|
|||
|
|
parser.add_argument('--fps', type=int, default=25, help='Video FPS (default: 25)')
|
|||
|
|
parser.add_argument('--max-frames', type=int, help='Maximum number of frames to process')
|
|||
|
|
parser.add_argument('--show-trajectory', action='store_true', help='Show trajectory lines')
|
|||
|
|
parser.add_argument('--compare', type=str, help='Second tracking.json for side-by-side comparison')
|
|||
|
|
parser.add_argument('--labels', type=str, nargs=2, default=['Method 1', 'Method 2'],
|
|||
|
|
help='Labels for comparison mode')
|
|||
|
|
parser.add_argument('--class-id', type=int, default=None,
|
|||
|
|
help='Filter by class ID (0=Car, 1=Pedestrian, 2=Cyclist, etc.). If not specified, all classes are shown.')
|
|||
|
|
parser.add_argument('--track-ids', type=int, nargs='+', default=None,
|
|||
|
|
help='Specific track IDs to visualize (e.g., --track-ids 1 2 3). If not specified, all tracks are shown.')
|
|||
|
|
parser.add_argument('--align-by-frame-id', action='store_true',
|
|||
|
|
help='Align tracking frames to the raw input video by frameId and export images only.')
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 检查输入文件
|
|||
|
|
tracking_path = Path(args.tracking)
|
|||
|
|
if not tracking_path.exists():
|
|||
|
|
print(f"Error: Tracking file not found: {tracking_path}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
image_source = Path(args.images)
|
|||
|
|
if not image_source.exists():
|
|||
|
|
print(f"Error: Image source not found: {image_source}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 加载跟踪数据
|
|||
|
|
print(f"Loading tracking data from {tracking_path}")
|
|||
|
|
with open(tracking_path, 'r') as f:
|
|||
|
|
tracking_data = json.load(f)
|
|||
|
|
|
|||
|
|
print(f"Loaded {len(tracking_data)} frames")
|
|||
|
|
|
|||
|
|
# 对比模式
|
|||
|
|
if args.compare:
|
|||
|
|
compare_path = Path(args.compare)
|
|||
|
|
if not compare_path.exists():
|
|||
|
|
print(f"Error: Comparison file not found: {compare_path}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
output_path = Path(args.output)
|
|||
|
|
if not str(output_path).endswith('.mp4'):
|
|||
|
|
output_path = output_path.with_suffix('.mp4')
|
|||
|
|
|
|||
|
|
create_side_by_side_comparison(
|
|||
|
|
tracking_path, compare_path, image_source,
|
|||
|
|
output_path, args.labels, args.max_frames or 100
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
# 单文件可视化
|
|||
|
|
# 准备过滤参数
|
|||
|
|
track_ids_set = set(args.track_ids) if args.track_ids else None
|
|||
|
|
|
|||
|
|
# 打印过滤信息
|
|||
|
|
if args.class_id is not None or args.track_ids:
|
|||
|
|
print(f"\nFiltering settings:")
|
|||
|
|
if args.class_id is not None:
|
|||
|
|
print(f" Class ID: {args.class_id}")
|
|||
|
|
if args.track_ids:
|
|||
|
|
print(f" Track IDs: {args.track_ids}")
|
|||
|
|
|
|||
|
|
visualize_tracking_on_images(
|
|||
|
|
tracking_data,
|
|||
|
|
image_source,
|
|||
|
|
args.output,
|
|||
|
|
save_video=args.save_video,
|
|||
|
|
fps=args.fps,
|
|||
|
|
max_frames=args.max_frames,
|
|||
|
|
show_trajectory=args.show_trajectory,
|
|||
|
|
class_id_filter=args.class_id,
|
|||
|
|
track_ids_filter=track_ids_set,
|
|||
|
|
align_by_frame_id=args.align_by_frame_id,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|