import json
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

json_file = 'semantic_map_2d.json'  # 你的2D语义json
occ_img = Image.open('test/839920/occupancy.png')
occ_arr = np.array(occ_img)
h, w = occ_arr.shape

with open(json_file) as f:
    instances = json.load(f)

# 构建类别到颜色的映射
all_labels = sorted(list(set([inst['category_label'] for inst in instances])))
cmap = plt.get_cmap('tab20', len(all_labels))
label2color = {label: cmap(i) for i, label in enumerate(all_labels)}

# 1. 先画背景occupancy
plt.figure(figsize=(12,12))
plt.imshow(occ_arr, cmap='gray', interpolation='none', alpha=0.7)

# 2. 遍历每个实例: 画mask（可选）、BBox
for inst in instances:
    mask_coords = inst['mask_coords']
    label = inst['category_label']
    bbox = inst['bbox']
    color = label2color[label]

    # [1] 画mask填充（可选注释掉），按类别着色，alpha半透明
    if mask_coords:
        ys, xs = zip(*mask_coords)
        plt.scatter(xs, ys, c=[color], s=8, marker='s', linewidths=0, alpha=0.7)

    # [2] 画边框
    xmin, ymin, xmax, ymax = bbox
    rect = plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, edgecolor=color, fill=False, linewidth=2, alpha=0.9)
    plt.gca().add_patch(rect)
    # [3] 打label文本
    plt.text(xmin+2, ymin+9, label, color=color, fontsize=9, va='top', ha='left',
             bbox=dict(facecolor='white',edgecolor='none',alpha=0.7,pad=1))

plt.title('Semantic 2D Map Overlay on Occupancy')
plt.axis('equal')
plt.xlim(0, w)
plt.ylim(h, 0)

# 图例
for label, color in label2color.items():
    plt.scatter([], [], c=[color], label=label)
plt.legend(fontsize=8,bbox_to_anchor=(1,1))
plt.tight_layout()
plt.show()