import json
import numpy as np
import matplotlib.pyplot as plt

# --- 选定分辨率（单位：米）
scale = 0.05

# --- 读取物理世界的2D mask信息
with open('2D_Semantic_Map_839920_Physic.json') as f:
    data = json.load(f)

# --- 自动推断物理世界mask的坐标边界
all_y = [float(y) for inst in data for y, x in inst['mask_coords_m']]
all_x = [float(x) for inst in data for y, x in inst['mask_coords_m']]
min_y, max_y = min(all_y), max(all_y)
min_x, max_x = min(all_x), max(all_x)
h = int(np.ceil((max_y - min_y) / scale)) + 1
w = int(np.ceil((max_x - min_x) / scale)) + 1

img = np.zeros((h, w), dtype=int)

# --- 物理世界 -> mask像素下标
def world2pix(x, y):
    """
    输入：x, y为浮点物理世界坐标(单位: 米)
    输出：py(img行), px(img列)
    """
    px = int(round((float(x) - min_x) / scale))
    py = int(round((float(y) - min_y) / scale))
    return py, px

# --- 遍历所有物体实例，填充mask
for inst in data:
    cid = inst['category_id']
    for y, x in inst['mask_coords_m']:
        py, px = world2pix(x, y)
        # 容错，防止列表范围越界
        if 0 <= py < h and 0 <= px < w:
            img[py, px] = cid

plt.figure(figsize=(12, 12))
img_extent = [min_x, min_x + w * scale, min_y, min_y + h * scale]
plt.imshow(img, cmap='tab20', extent=img_extent, origin='lower')

# --- 画bbox（均已是物理世界坐标，直接用）
for inst in data:
    bbox = inst['bbox_m']
    x0, y0, x1, y1 = map(float, bbox)
    plt.gca().add_patch(
        plt.Rectangle((x0, y0 + 0.5 * scale), x1 - x0, y1 - y0,
                      edgecolor='red', fill=False, linewidth=1)
    )
plt.title('2D Semantic Map (Restored from Physical JSON)')
plt.xlabel('X (meters)')
plt.ylabel('Y (meters)')
plt.savefig('2D_Semantic_Map_839920_Physic_Restored.png', bbox_inches='tight', dpi=300)
plt.show()