hongchen/door_detect/analyze_poses.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
分析全景图点位，找出最靠近入户门的点位
- 基于拍摄顺序（id 递增）和可见性关系
- 将 3D 世界坐标映射到 2D 平面可视化
- 用红色标记最靠近入户门的点，绿色标记其他点
"""

import json
import os

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle


def load_vision_data(filepath):
    """加载 vision.txt 数据"""
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
        # 找到 JSON 开始的位置
        start_idx = content.find('{')
        data = json.loads(content[start_idx:])
    return data


def extract_poses(data):
    """提取所有点位的 pose 信息"""
    poses = []
    for sweep in data['sweepLocations']:
        pose_info = {
            'id': sweep['id'],
            'uuid': sweep['uuid'],
            'x': sweep['pose']['translation']['x'],
            'y': sweep['pose']['translation']['y'],
            'z': sweep['pose']['translation']['z'],
            'visibles': sweep.get('visibles', []),
            'visibles2': sweep.get('visibles2', []),
            'visibles3': sweep.get('visibles3', [])
        }
        poses.append(pose_info)
    return poses


def find_entrance_point(poses):
    """
    找出最靠近入户门的点位

    逻辑：
    1. 入户门位于边缘（几何中心最远的点）
    2. 入户门在走廊上（可见性符合距离规律：近的能看见，远的看不见）
    3. 房间点会被排除（近的点可能看不见，因为有墙遮挡）

    关键洞察：
    - 走廊是开阔的，可见性主要由距离决定
    - 房间内有墙体，可见性会被墙阻挡
    - 如果"近的点看不见"但"远的点反而能看见"，说明在房间内
    """

    # 计算几何中心
    center_x = sum(p['x'] for p in poses) / len(poses)
    center_y = sum(p['y'] for p in poses) / len(poses)

    # 计算所有点对之间的距离
    n = len(poses)
    distances = {}  # distances[(id1, id2)] = 距离
    for i in range(n):
        for j in range(i + 1, n):
            p1, p2 = poses[i], poses[j]
            dx = p1['x'] - p2['x']
            dy = p1['y'] - p2['y']
            dist = np.sqrt(dx ** 2 + dy ** 2)
            distances[(p1['id'], p2['id'])] = dist
            distances[(p2['id'], p1['id'])] = dist

    # 对每个点，分析其可见性与距离的关系
    corridor_scores = {}  # 走廊分数（越高越可能在走廊）

    for p in poses:
        pid = p['id']
        visibles = set(p['visibles'])

        # 收集该点到其他所有点的距离和可见性
        visible_distances = []  # 可见的点的距离
        invisible_distances = []  # 不可见的点的距离

        for other in poses:
            if other['id'] == pid:
                continue
            dist = distances[(pid, other['id'])]
            if other['id'] in visibles:
                visible_distances.append(dist)
            else:
                invisible_distances.append(dist)

        # 关键分析：
        # 1. 走廊点：可见的最大距离 > 不可见的最小距离（远的能看见，近的肯定能看见）
        # 2. 房间点：存在不可见的近点，但存在可见的远点

        if len(invisible_distances) == 0:
            # 所有点都可见，肯定是走廊中心点
            corridor_scores[pid] = 1.0
            continue

        min_invisible_dist = min(invisible_distances)  # 最近的不可见点
        max_visible_dist = max(visible_distances) if visible_distances else 0  # 最远的可见点

        # 计算"走廊分数"
        # 如果 min_invisible_dist > max_visible_dist，说明可见性符合距离规律（走廊）
        # 如果 min_invisible_dist < max_visible_dist，说明有异常（房间）

        if max_visible_dist == 0:
            corridor_scores[pid] = 0.0  # 什么都看不见，可能在死角
        else:
            ratio = min_invisible_dist / max_visible_dist
            # ratio > 1: 走廊（所有不可见点都比可见点远）
            # ratio < 1: 房间（存在不可见的近点）
            corridor_scores[pid] = ratio / (1 + ratio)  # 映射到 0-1

    # 计算边缘度（距离中心的距离）
    center_distances = {}
    for p in poses:
        dx = p['x'] - center_x
        dy = p['y'] - center_y
        center_distances[p['id']] = np.sqrt(dx ** 2 + dy ** 2)
    max_center_dist = max(center_distances.values())

    # 综合分数：走廊分数 * 边缘度
    entrance_scores = {}
    for p in poses:
        pid = p['id']
        edge_score = center_distances[pid] / max_center_dist
        corridor_score = corridor_scores[pid]
        # 入户门 = 边缘 + 走廊
        entrance_scores[pid] = edge_score * 0.6 + corridor_score * 0.4

    # 输出分析结果
    print("=== 几何中心 ===")
    print(f"中心坐标：({center_x:.3f}, {center_y:.3f})")

    print("\n=== 点位分析（可见性与距离关系）===")
    for pid in sorted(corridor_scores.keys()):
        p = next(po for po in poses if po['id'] == pid)
        visibles = p['visibles']
        visible_dists = [distances[(pid, vid)] for vid in visibles if (pid, vid) in distances]
        invisible_dists = [distances[(pid, oid)] for oid in
                           [po['id'] for po in poses if po['id'] != pid and po['id'] not in visibles]]

        min_inv = min(invisible_dists) if invisible_dists else float('inf')
        max_vis = max(visible_dists) if visible_dists else 0

        print(f"点 {pid}: 可见{len(visibles)}个点，" +
              f"最近不可见={min_inv:.2f}m, 最远可见={max_vis:.2f}m, " +
              f"走廊分数={corridor_scores[pid]:.3f}")

    print("\n=== 入口分数（边缘度 60% + 走廊分数 40%）===")
    for pid in sorted(entrance_scores.keys(), key=lambda x: entrance_scores[x], reverse=True):
        print(f"点 {pid}: 分数={entrance_scores[pid]:.3f}")

    entrance_id = max(entrance_scores, key=entrance_scores.get)
    print(f"\n>>> 最靠近入户门的点位：ID = {entrance_id}")

    return entrance_id


def transform_to_world_coords(poses):
    """
    将所有拍摄点位通过位姿转换到世界坐标系

    对于每个 sweep，其 pose.translation 已经是世界坐标
    这里统一以第一个点 (id=0) 作为参考原点，将所有点转换到局部坐标系

    Args:
        poses: 包含位姿信息的列表

    Returns:
        转换后的点位列表
    """
    # 以第一个点作为坐标原点
    origin_pose = poses[0]
    origin_x = origin_pose['x']
    origin_y = origin_pose['y']
    origin_z = origin_pose['z']

    # 获取第一个点的旋转（用于方向对齐）
    origin_rot = origin_pose.get_rotation() if hasattr(origin_pose, 'get_rotation') else None

    transformed_poses = []
    for p in poses:
        # 平移变换：将原点移到第一个点
        world_x = p['x'] - origin_x
        world_y = p['y'] - origin_y
        world_z = p['z'] - origin_z  # 保持高度信息

        transformed_poses.append({
            'id': p['id'],
            'x': world_x,
            'y': world_y,
            'z': world_z,
            'visibles': p['visibles'],
            'original_x': p['x'],
            'original_y': p['y'],
            'original_z': p['z']
        })

    print(f"\n=== 坐标变换 ===")
    print(f"参考原点 (ID=0): ({origin_x:.3f}, {origin_y:.3f}, {origin_z:.3f})")
    print(f"变换后所有点位偏移了：(-{origin_x:.3f}, -{origin_y:.3f}, -{origin_z:.3f})")

    return transformed_poses


def project_to_2d(transformed_poses):
    """
    将世界坐标系下的 3D 点映射到 XY 平面（俯视图）
    忽略 Z 坐标，只保留 X, Y 用于平面绘制

    根据 demo.png，直接使用原始坐标，不交换轴
    """
    points_2d = []
    for p in transformed_poses:
        points_2d.append({
            'id': p['id'],
            'x': p['x'],  # 世界坐标系 X（横向）
            'y': p['y'],  # 世界坐标系 Y（纵向）
            'z': p['z'],  # 高度（用于参考）
            'visibles': p['visibles']
        })
    return points_2d


def visualize_poses(poses_2d, entrance_id, output_path):
    """可视化点位分布"""

    fig, ax = plt.subplots(figsize=(12, 10))

    # 提取所有坐标
    xs = [p['x'] for p in poses_2d]
    ys = [p['y'] for p in poses_2d]

    # 计算坐标范围
    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)

    # 设置坐标范围（带边距）
    x_margin = (x_max - x_min) * 0.1
    y_margin = (y_max - y_min) * 0.1
    ax.set_xlim(x_min - x_margin, x_max + x_margin)
    ax.set_ylim(y_min - y_margin, y_max + y_margin)

    # 保持 Y 轴方向与世界坐标系一致（上北下南）
    # ax.invert_yaxis()  # 如果需要俯视效果可取消注释

    # 绘制连线（可见性关系）
    for p in poses_2d:
        for vid in p['visibles']:
            # 找到对应的点
            target = next((pt for pt in poses_2d if pt['id'] == vid), None)
            if target:
                ax.plot([p['x'], target['x']], [p['y'], target['y']],
                        'gray', linewidth=0.5, alpha=0.3)

    # 绘制点位
    for p in poses_2d:
        if p['id'] == entrance_id:
            # 入口点 - 红色实心圆
            circle = Circle((p['x'], p['y']), radius=0.3,
                            color='red', fill=True,
                            label='入口点' if p['id'] == entrance_id else '')
            ax.add_patch(circle)
            # 添加 ID 标签
            ax.annotate(f'ID:{p["id"]}', (p['x'], p['y']),
                        textcoords="offset points", xytext=(0, 15),
                        ha='center', color='red', fontsize=12, fontweight='bold')
        else:
            # 其他点 - 绿色实心圆
            circle = Circle((p['x'], p['y']), radius=0.2,
                            color='green', fill=True)
            ax.add_patch(circle)
            # 添加 ID 标签
            ax.annotate(f'ID:{p["id"]}', (p['x'], p['y']),
                        textcoords="offset points", xytext=(0, -12),
                        ha='center', color='green', fontsize=9)

    # 添加图例
    red_circle = plt.Line2D([], [], marker='o', markersize=15,
                            markerfacecolor='red', markeredgewidth=0,
                            label='Entrance Point')
    green_circle = plt.Line2D([], [], marker='o', markersize=12,
                              markerfacecolor='green', markeredgewidth=0,
                              label='Other Points')
    ax.legend(handles=[red_circle, green_circle], loc='upper right')

    # 设置标签和标题
    ax.set_xlabel('X (meter)', fontsize=11)
    ax.set_ylabel('Y (meter)', fontsize=11)
    ax.set_title('Panoramic Point Distribution\n(Red=Near Entrance, Green=Other Points)', fontsize=13)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\n可视化结果已保存至：{output_path}")

    return fig, ax


def main(vision_file, output_image):
    print("加载数据...")
    data = load_vision_data(vision_file)

    print("提取点位信息...")
    poses = extract_poses(data)
    print(f"共 {len(poses)} 个点位")

    # 打印每个点位的详细信息
    print("\n=== 点位详情 ===")
    for p in poses:
        print(f"ID {p['id']:2d}: pos=({p['x']:8.3f}, {p['y']:8.3f}, {p['z']:8.3f}), visibles={p['visibles']}")

    # 找出入口点
    print("\n=== 分析入户门位置 ===")
    entrance_id = find_entrance_point(poses)

    # 投影到 2D
    poses_2d = project_to_2d(poses)

    # 可视化
    visualize_poses(poses_2d, entrance_id, output_image)

    # 打印结论
    print("\n" + "=" * 50)
    print("结 论:")
    print("=" * 50)
    entrance_point = next(p for p in poses if p['id'] == entrance_id)
    print(f"最靠近入户门的点位：ID = {entrance_id}")
    print(f"  世界坐标：({entrance_point['x']:.3f}, {entrance_point['y']:.3f}, {entrance_point['z']:.3f})")
    print(f"  可见点位：{entrance_point['visibles']}")
    print(f"\n其余点位 (绿色)：{[p['id'] for p in poses if p['id'] != entrance_id]}")


if __name__ == '__main__':
    vision_folder = "/home/gu/视频/door_detect/high_temp/ten_pano_demo/depth"
    out_folder = "/home/gu/视频/door_detect/high_temp/ten_pano_demo/2d_image"
    vision_scenes = os.listdir(vision_folder)
    for scene in vision_scenes:
        scene_name = os.path.join(vision_folder, scene, 'vision.txt')
        out_name = os.path.join(out_folder, scene + ".png")
        main(scene_name, out_name)