#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 入户门投影可视化 将所有检测到的门映射到 2D 俯视图上： - 绿色圆点：普通门（从 JSON 中的 doors 提取） - 红色圆点：入户门（entrance_door_id 对应的门 或 world_position） """ import json import os import numpy as np import matplotlib.pyplot as plt from matplotlib.patches import Circle from pathlib import Path def load_entrance_json(json_path: str) -> dict: """加载入户门 JSON 文件""" with open(json_path, 'r', encoding='utf-8') as f: return json.load(f) def visualize_doors_2d( entrance_data: dict, output_path: str ): """ 在 2D 俯视图上绘制所有门和入户门 Args: entrance_data: 入户门 JSON 数据 output_path: 输出图片路径 """ scene_name = entrance_data.get('scene_name', 'unknown') entrance_door_id = entrance_data.get('entrance_door_id') world_position = entrance_data.get('world_position') doors = entrance_data.get('doors', []) print(f"\n处理场景：{scene_name}") print(f" entrance_door_id: {entrance_door_id}") print(f" world_position: {world_position}") print(f" 门数量：{len(doors)}") # 确定入户门的世界坐标 entrance_world_center = None entrance_door_obj = None if entrance_door_id is not None: # 从 doors 中找到对应的门 for door in doors: if door.get('door_id') == entrance_door_id: entrance_world_center = door.get('world_center') entrance_door_obj = door print(f" 使用 entrance_door_id={entrance_door_id} 的世界坐标") break if entrance_world_center is None and world_position is not None: entrance_world_center = world_position print(f" 使用 world_position 作为入户门位置") if entrance_world_center is None: print(" ⚠️ 警告：无法确定入户门世界坐标，只绘制绿色门位置") # 提取所有门的中心坐标（2D 俯视，使用 X-Y 平面） # world_center = [x, y, z]，其中 Y 是前后方向，X 是左右方向 door_centers = [] for door in doors: center = door.get('world_center') if center: door_centers.append({ 'door_id': door.get('door_id'), 'x': center[0], # X 坐标（左右） 'y': center[1], # Y 坐标（前后） 'z': center[2], # Z 坐标（高度） 'is_entrance': door.get('door_id') == entrance_door_id, 'source_image': door.get('source_image'), 'confidence': door.get('confidence'), 'size': door.get('bounding_box_size', {}) }) if not door_centers: print(" ⚠️ 没有有效的门坐标") return # 创建图表 fig, ax = plt.subplots(figsize=(14, 12)) # 提取坐标 xs = [p['x'] for p in door_centers] ys = [p['y'] for p in door_centers] # 计算坐标范围 x_min, x_max = min(xs), max(xs) y_min, y_max = min(ys), max(ys) # 设置坐标范围（带边距） x_margin = max((x_max - x_min) * 0.15, 1.0) y_margin = max((y_max - y_min) * 0.15, 1.0) ax.set_xlim(x_min - x_margin, x_max + x_margin) ax.set_ylim(y_min - y_margin, y_max + y_margin) # 绘制所有门 for door in door_centers: is_entrance = (door['door_id'] == entrance_door_id) or \ (entrance_world_center and np.allclose([door['x'], door['z']], [entrance_world_center[0], entrance_world_center[2]], atol=0.1)) if is_entrance: # 入户门 - 红色实心圆（大） circle = Circle((door['x'], door['y']), radius=0.4, facecolor='red', edgecolor='darkred', linewidth=2) ax.add_patch(circle) # 添加标签 label = f"ID:{door['door_id']} (Entrance)" ax.annotate(label, (door['x'], door['y']), textcoords="offset points", xytext=(0, 20), ha='center', color='red', fontsize=11, fontweight='bold') else: # 普通门 - 绿色实心圆（小） circle = Circle((door['x'], door['y']), radius=0.25, facecolor='green', edgecolor='darkgreen', linewidth=1.5) ax.add_patch(circle) # 添加标签 label = f"ID:{door['door_id']}" ax.annotate(label, (door['x'], door['y']), textcoords="offset points", xytext=(0, -15), ha='center', color='green', fontsize=9) # 如果只有 world_position 而没有 entrance_door_id，绘制一个蓝色叉号标记 if entrance_door_id is None and world_position is not None: ax.plot(world_position[0], world_position[1], 'bx', markersize=20, markeredgewidth=3, label='World Position (Estimated)') # 添加图例 red_circle = plt.Line2D([], [], marker='o', markersize=18, markerfacecolor='red', markeredgewidth=2, markeredgecolor='darkred', label='Entrance Door') green_circle = plt.Line2D([], [], marker='o', markersize=14, markerfacecolor='green', markeredgewidth=1.5, markeredgecolor='darkgreen', label='Regular Door') ax.legend(handles=[red_circle, green_circle], loc='upper right', fontsize=11) # 设置标签和标题 ax.set_xlabel('X (meter)', fontsize=12) ax.set_ylabel('Y (meter)', fontsize=12) ax.set_title(f'Scene: {scene_name}\nDoor Distribution (Red=Entrance, Green=Regular)', fontsize=14) ax.set_aspect('equal') ax.grid(True, alpha=0.3) # 添加信息框 info_text = f"Total Doors: {len(doors)}\n" if entrance_door_id is not None: info_text += f"Entrance Door ID: {entrance_door_id}" elif world_position is not None: info_text += f"Entrance (from world_position)" else: info_text += "Entrance: Not specified" props = dict(boxstyle='round', facecolor='wheat', alpha=0.5) ax.text(0.02, 0.98, info_text, transform=ax.transAxes, fontsize=10, verticalalignment='top', bbox=props) plt.tight_layout() # 创建输出目录 os.makedirs(os.path.dirname(output_path), exist_ok=True) plt.savefig(output_path, dpi=150, bbox_inches='tight') print(f" 可视化结果已保存：{output_path}") return fig, ax def process_all_scenes( entrance_json_base: str, output_base: str ): """ 处理所有场景 Args: entrance_json_base: entrance_json_output 基础目录 output_base: 输出基础目录 """ json_files = list(Path(entrance_json_base).glob("*.json")) print(f"找到 {len(json_files)} 个场景文件") for json_file in json_files: scene_name = json_file.stem output_path = Path(output_base) / f"{scene_name}_doors.png" entrance_data = load_entrance_json(str(json_file)) visualize_doors_2d(entrance_data, str(output_path)) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( description="入户门 2D 可视化（俯视图）", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # 处理单个场景 python visualize_entrance_2d.py --scene KJ-dg21mhHcJLW # 处理所有场景 python visualize_entrance_2d.py --all # 指定自定义路径 python visualize_entrance_2d.py --all \\ --json-base /path/to/entrance_json_output \\ --output-base /path/to/output """ ) parser.add_argument("--scene", type=str, help="处理单个场景名称") parser.add_argument("--all", action="store_true", help="处理所有场景") # 路径参数 parser.add_argument("--json-base", type=str, default="/home/gu/视频/door_detect/high_temp/ten_pano_demo/entrance_json_output", help="entrance_json_output 基础目录") parser.add_argument("--output-base", type=str, default="/home/gu/视频/door_detect/high_temp/ten_pano_demo/entrance_vis_2d", help="输出基础目录") args = parser.parse_args() if args.scene: # 处理单个场景 json_file = Path(args.json_base) / f"{args.scene}.json" if not json_file.exists(): print(f"❌ JSON 文件不存在：{json_file}") else: output_path = Path(args.output_base) / f"{args.scene}_doors.png" entrance_data = load_entrance_json(str(json_file)) visualize_doors_2d(entrance_data, str(output_path)) elif args.all: # 处理所有场景 process_all_scenes( entrance_json_base=args.json_base, output_base=args.output_base ) else: parser.print_help()