src/visualize_footprint_json.py

import json
import cv2
import numpy as np
from pathlib import Path

def visualize_footprint_json(json_path, output_dir):
    """可视化足印JSON数据"""
    # 读取JSON数据
    with open(json_path, 'r') as f:
        data = json.load(f)
    
    # 创建输出目录
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 确定画布大小（从数据中获取最大坐标）
    max_x = max_y = 0
    for frame in data['frames']:
        for footprint in frame['footprints']:
            pos = footprint['position']
            max_x = max(max_x, pos['x'] + pos['width'])
            max_y = max(max_y, pos['y'] + pos['height'])
    
    for area in data['footprintArea']:
        pos = area['areaPosition']
        max_x = max(max_x, pos['x'] + pos['width'])
        max_y = max(max_y, pos['y'] + pos['height'])
    
    # 添加一些边距
    canvas_width = max_x + 50
    canvas_height = max_y + 50
    
    # 颜色映射
    color_map = {
        'LF': (0, 255, 0),      # 绿色（原leftFront）
        'RF': (0, 255, 255),    # 黄色（原rightFront）
        'LH': (255, 0, 255),    # 紫色（原leftHind）
        'RH': (0, 165, 255),    # 橙色（原rightHind）
        'unknown': (128, 128, 128)
    }
    
    # 1. 绘制所有足印区域的覆盖范围
    area_canvas = np.ones((canvas_height, canvas_width, 3), dtype=np.uint8) * 255
    for area in data['footprintArea']:
        pos = area['areaPosition']
        color = color_map[area['type']]
        
        # 绘制区域边界框
        cv2.rectangle(area_canvas, 
                     (pos['x'], pos['y']), 
                     (pos['x'] + pos['width'], pos['y'] + pos['height']), 
                     color, 2)
        
        # 添加标签
        label = f"{area['type']}#{area['footprintAreaId'].split('_')[-1]}"
        cv2.putText(area_canvas, label, 
                    (pos['x'], pos['y'] - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    
    # 保存足印区域图
    cv2.imwrite(str(output_dir / 'footprint_areas.png'), area_canvas)
    
    # 2. 为每一帧创建可视化
    frame_ids = sorted(list(set(frame['frameId'] for frame in data['frames'])))
    for frame_id in frame_ids:
        frame_canvas = np.ones((canvas_height, canvas_width, 3), dtype=np.uint8) * 255
        
        # 首先绘制所有区域的轮廓（半透明）
        overlay = frame_canvas.copy()
        for area in data['footprintArea']:
            if area['startFrame'] <= frame_id <= area['endFrame']:
                pos = area['areaPosition']
                color = color_map[area['type']]
                cv2.rectangle(overlay, 
                            (pos['x'], pos['y']), 
                            (pos['x'] + pos['width'], pos['y'] + pos['height']), 
                            color, -1)  # 填充矩形
        
        # 应用半透明效果
        alpha = 0.3
        cv2.addWeighted(overlay, alpha, frame_canvas, 1 - alpha, 0, frame_canvas)
        
        # 然后绘制当前帧的足印
        for frame_data in data['frames']:
            if frame_data['frameId'] == frame_id:
                for footprint in frame_data['footprints']:
                    pos = footprint['position']
                    color = color_map[footprint['type']]
                    
                    # 绘制足印边界框
                    cv2.rectangle(frame_canvas, 
                                (pos['x'], pos['y']), 
                                (pos['x'] + pos['width'], pos['y'] + pos['height']), 
                                color, 2)
                    
                    # 添加标签
                    label = f"{footprint['type']}#{footprint['id'].split('_')[-1]}"
                    cv2.putText(frame_canvas, label, 
                              (pos['x'], pos['y'] - 5),
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
        
        # 添加帧号
        cv2.putText(frame_canvas, f"Frame: {frame_id}", 
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
        
        # 保存帧图像
        cv2.imwrite(str(output_dir / f'frame_{frame_id:04d}.png'), frame_canvas)
    
    print(f"可视化结果已保存至: {output_dir}")

def main():
    # 假设JSON文件在results目录下的最新时间戳文件夹中
    results_dir = Path("results")
    latest_dir = max(results_dir.glob("*"), key=lambda p: p.stat().st_mtime)
    json_path = latest_dir / "data" / "footprint_data.json"
    
    # 创建可视化输出目录
    output_dir = latest_dir / "visualization"
    
    visualize_footprint_json(json_path, output_dir)

if __name__ == "__main__":
    main()