hongchen/Floorplan/vis.py

import json
import cv2
import numpy as np
import os
from collections import defaultdict

# ================= 1. 颜色映射 =================

COLOR_MAP = {
    "living_room": (255, 180, 100),  # BGR: 橙色
    "bed_room": (100, 255, 100),  # BGR: 绿色
    "bath_room": (255, 100, 255),  # BGR: 紫色
    "kitchen_room": (100, 255, 255),  # BGR: 黄色
    "other_room": (180, 180, 180),  # BGR: 灰色
    "balcony": (200, 150, 100),  # BGR: 棕色
}

DEFAULT_COLOR = (200, 200, 200)


# ================= 2. 工具函数 =================

def get_image_size(data):
    """从 JSON 数据获取图像尺寸"""
    if 'image_size' in data:
        return data['image_size']['width'], data['image_size']['height']

    # 从线段数据推断
    max_x, max_y = 0, 0
    for block in data.get('block', []):
        segments = block.get('points', [])
        if isinstance(segments, list) and len(segments) > 0:
            if isinstance(segments[0], list) and len(segments[0]) == 4:
                for seg in segments:
                    max_x = max(max_x, seg[0], seg[2])
                    max_y = max(max_y, seg[1], seg[3])

    for area in data.get('connect_area', []):
        max_x = max(max_x, area.get('x', 0) + area.get('w', 0))
        max_y = max(max_y, area.get('y', 0) + area.get('h', 0))

    return int(max_x) + 100, int(max_y) + 100


def get_points_from_segments(segments):
    """从线段列表提取所有唯一点"""
    if not segments:
        return []

    points_set = set()
    for seg in segments:
        if len(seg) == 4:
            points_set.add((int(seg[0]), int(seg[1])))
            points_set.add((int(seg[2]), int(seg[3])))

    return list(points_set)


# ================= 3. 主可视化函数 =================

def visualize_final_json(json_path, output_path, rgb_path=None, vis_door=False):
    """
    最终可视化：线段格式 JSON

    修改点：
    1. 房间端点：白色点 (255, 255, 255)
    2. 标签位置：使用 JSON 中的 center 字段
    """
    print(f"\n{'=' * 60}")
    print(f"🎨 最终可视化：{os.path.basename(json_path)}")
    print('=' * 60)

    # 1. 读取 JSON
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 2. 获取尺寸并创建画布
    width, height = get_image_size(data)

    if rgb_path and os.path.exists(rgb_path):
        canvas = cv2.imread(rgb_path)
        print(f"🖼️ 使用背景图：{rgb_path}")
    else:
        # canvas = np.ones((height, width, 3), dtype=np.uint8) * 255
        canvas = np.zeros((height, width, 3), dtype=np.uint8)
        print(f"📐 创建白色画布：{width}x{height}")

    # 3. 绘制房间块 (Blocks)
    blocks = data.get('block', [])
    print(f"\n🏠 绘制 {len(blocks)} 个房间块...")

    for block_idx, block in enumerate(blocks):
        block_id = block.get('id', block_idx)
        label = block.get('label', 'door')
        segments = block.get('points', [])
        center = block.get('center', None)  # 【修改】使用 JSON 中的 center 字段

        # 获取颜色
        color = COLOR_MAP.get(label.lower(), DEFAULT_COLOR)

        # 验证线段格式
        if not (isinstance(segments, list) and len(segments) > 0):
            print(f"  ⚠️ Block {block_id}: 无数据，跳过")
            continue

        if not (isinstance(segments[0], list) and len(segments[0]) == 4):
            print(f"  ⚠️ Block {block_id}: 格式不正确，跳过")
            continue

        # 提取所有点
        points = get_points_from_segments(segments)

        if len(points) < 2:
            print(f"  ⚠️ Block {block_id}: 点数不足，跳过")
            continue

        # 【1】绘制线段（按元素连线）
        for seg in segments:
            p1 = (int(seg[0]), int(seg[1]))
            p2 = (int(seg[2]), int(seg[3]))
            cv2.line(canvas, p1, p2, color, 2, cv2.LINE_AA)

        # 【2】绘制点（所有顶点标记）【修改】白色点
        for pt in points:
            cv2.circle(canvas, pt, 4, (255, 255, 255), -1)  # 白色点

        # 【3】绘制标签【修改】使用 center 字段
        if center and len(center) == 2:
            cx, cy = int(center[0]), int(center[1])
        else:
            # 如果没有 center，回退到计算中心
            pts = np.array(points, dtype=np.int32)
            M = cv2.moments(pts)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
            else:
                cx = int(np.mean(pts[:, 0]))
                cy = int(np.mean(pts[:, 1]))
            print(f"  ⚠️ Block {block_id}: 无 center 字段，使用计算中心")

        # 确保在画布内
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.6
        thickness = 1
        (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)

        cx = max(tw // 2 + 5, min(cx, width - tw // 2 - 5))
        cy = max(th // 2 + 5, min(cy, height - th // 2 - 5))

        # 标签背景（白色矩形）
        cv2.rectangle(canvas, (cx - tw // 2 - 3, cy - th - 3),
                      (cx + tw // 2 + 3, cy + 3), (255, 255, 255), -1)

        # 标签文字（黑色）
        cv2.putText(canvas, label, (cx - tw // 2, cy),
                    font, font_scale, (0, 0, 0), thickness)

        # Block ID（灰色小字）
        id_text = f"ID:{block_id}"
        (iw, ih), _ = cv2.getTextSize(id_text, font, 0.4, 1)
        cv2.putText(canvas, id_text, (cx - iw // 2, cy + 15),
                    font, 0.4, (100, 100, 100), 1)

        print(f"  ✅ Block {block_id} ({label}): {len(segments)} 线段，{len(points)} 点，center=({cx},{cy})")

    # 4. 绘制连接区域 (Connect Area)
    connect_areas = data.get('connect_area', [])
    print(f"\n🔗 绘制 {len(connect_areas)} 个连接区域...")

    for area_idx, area in enumerate(connect_areas):
        area_id = area.get('id', area_idx)
        x = area.get('x', 0)
        y = area.get('y', 0)
        w = area.get('w', 0)
        h = area.get('h', 0)
        label = area.get('label', 'door').lower()
        block_pair = area.get('block_pair', [])

        # 确保在画布内
        x = max(0, min(x, width - 1))
        y = max(0, min(y, height - 1))
        w = max(1, min(w, width - x))
        h = max(1, min(h, height - y))

        # 【unknown/unknow】红色实心矩形
        if vis_door and ("door" in label):
            cv2.rectangle(canvas, (x, y), (x + w, y + h), (0, 0, 255), -1)  # 红色填充
            cv2.rectangle(canvas, (x, y), (x + w, y + h), (255, 255, 255), 1)  # 白色边框
            print(f"  ✅ Connect {area_id}: 🔴 红色实心矩形")

    # 5. 绘制家具 (Furniture)
    furniture_list = data.get('furniture', [])
    print(f"\n🛋️ 绘制 {len(furniture_list)} 个家具...")

    FURNITURE_COLOR = (0, 165, 255)  # BGR: 橙色
    for item in furniture_list:
        label = item.get('label', '')
        center = item.get('center', None)
        pts = item.get('points', {})
        if pts:
            bx1, by1 = pts['x1'], pts['y1']
            bx2, by2 = pts['x3'], pts['y3']
            cv2.rectangle(canvas, (bx1, by1), (bx2, by2), FURNITURE_COLOR, 2)
        if center and len(center) == 2:
            cx, cy = int(center[0]), int(center[1])
            font = cv2.FONT_HERSHEY_SIMPLEX
            (tw, th), _ = cv2.getTextSize(label, font, 0.5, 1)
            cv2.rectangle(canvas, (cx - tw // 2 - 3, cy - th - 3),
                          (cx + tw // 2 + 3, cy + 3), FURNITURE_COLOR, -1)
            cv2.putText(canvas, label, (cx - tw // 2, cy),
                        font, 0.5, (255, 255, 255), 1)
        print(f"  ✅ {label} center=({center})")

    # 6. 添加图例
    print(f"\n📋 添加图例...")
    legend_y = 30
    legend_x = 20

    cv2.putText(canvas, "Legend:", (legend_x, legend_y - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    legend_y += 5

    for room_type, color in COLOR_MAP.items():
        if room_type in ["door"]:
            continue
        cv2.rectangle(canvas, (legend_x, legend_y),
                      (legend_x + 20, legend_y + 15), color, -1)
        cv2.rectangle(canvas, (legend_x, legend_y),
                      (legend_x + 20, legend_y + 15), (0, 0, 0), 1)
        cv2.putText(canvas, room_type, (legend_x + 25, legend_y + 12),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)
        legend_y += 20

    # 添加 unknown 图例
    cv2.rectangle(canvas, (legend_x, legend_y),
                  (legend_x + 20, legend_y + 15), (0, 0, 255), -1)
    cv2.rectangle(canvas, (legend_x, legend_y),
                  (legend_x + 20, legend_y + 15), (0, 0, 0), 1)
    cv2.putText(canvas, "door", (legend_x + 25, legend_y + 12),
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)

    # 6. 保存结果
    cv2.imwrite(output_path, canvas)

    print(f"\n{'=' * 60}")
    print(f"✨ 可视化完成！")
    print(f"   📁 输出路径：{output_path}")
    print(f"   📐 画布尺寸：{width}x{height}")
    print(f"   🏠 Block 数量：{len(blocks)}")
    print(f"   🔗 Connect Area 数量：{len(connect_areas)}")
    print('=' * 60)

    return canvas


# ================= 4. 批量处理 =================

def visualize_batch(json_folder, output_folder, rgb_folder=None, vis_door=False):
    """批量可视化文件夹中的所有 JSON"""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    json_files = [f for f in os.listdir(json_folder) if f.endswith('.json')]

    print(f"\n📁 发现 {len(json_files)} 个 JSON 文件\n")

    for json_file in json_files:
        json_path = os.path.join(json_folder, json_file)
        output_path = os.path.join(output_folder, json_file.replace('.json', '.png'))

        rgb_path = None
        if rgb_folder:
            rgb_path = os.path.join(rgb_folder, json_file.replace('.json', '.png'))
            if not os.path.exists(rgb_path):
                rgb_path = None

        try:
            visualize_final_json(json_path, output_path, rgb_path, vis_door)
        except Exception as e:
            print(f"❌ 处理 {json_file} 失败：{e}")


# ================= 5. 主运行流程 =================

if __name__ == '__main__':
    import sys

    if len(sys.argv) < 2:
        print("用法：python vis.py <文件夹名称>")
        print("示例：python vis.py SG-n6nV8B2oW95")
        sys.exit(1)

    folder_name = sys.argv[1]
    img_folder = "temp_data"
    folder_path = os.path.join(img_folder, folder_name)

    if not os.path.isdir(folder_path):
        print(f"错误：文件夹不存在 {folder_path}")
        sys.exit(1)

    # 查找对应的 JSON 文件（文件名与文件夹同名）
    json_files = [f for f in os.listdir(folder_path)
                  if f.startswith(folder_name) and f.endswith('.json')]

    if not json_files:
        print(f"错误：未找到 JSON 文件")
        sys.exit(1)

    json_name = json_files[0]
    json_path = os.path.join(folder_path, json_name)

    # 查找对应的 RGB 图片
    rgb_name = json_name.replace('.json', '.png')
    rgb_path = os.path.join(folder_path, rgb_name)

    if not os.path.exists(rgb_path):
        print(f"错误：未找到 {rgb_name}")
        sys.exit(1)

    # 输出可视化路径：name_vis.png（保存到子文件夹）
    vis_name = json_name.replace('.json', '_vis.png')
    output_path = os.path.join(folder_path, vis_name)

    print(f"处理文件夹：{folder_name}")
    print(f"   JSON: {json_path}")
    print(f"   RGB: {rgb_path}")
    print(f"   Vis: {output_path}")

    try:
        visualize_final_json(json_path, output_path, rgb_path, vis_door=True)
        print(f"✅ 完成 -> {os.path.basename(output_path)}")
    except Exception as e:
        print(f"❌ 错误：{e}")
        sys.exit(1)