hongchen/door_detect/visualize_entrance_2d.py

#!/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()