""" 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()