(img_color, labels, img_size, save_path)
| 231 | |
| 232 | |
| 233 | def visualize_2d(img_color, labels, img_size, save_path): |
| 234 | import matplotlib.pyplot as plt |
| 235 | # from skimage.segmentation import mark_boundaries |
| 236 | # from skimage.color import label2rgb |
| 237 | label_names = ["wall", "floor", "cabinet", "bed", "chair", |
| 238 | "sofa", "table", "door", "window", "bookshelf", |
| 239 | "picture", "counter", "desk", "curtain", "refridgerator", |
| 240 | "shower curtain", "toilet", "sink", "bathtub", "other"] |
| 241 | colors = np.array(list(SCANNET_COLOR_MAP_20.values()))[1:] |
| 242 | segmentation_color = np.zeros((img_size[0], img_size[1], 3)) |
| 243 | for i, color in enumerate(colors): |
| 244 | segmentation_color[labels == i] = color |
| 245 | alpha = 1 |
| 246 | overlay = (img_color * (1-alpha) + segmentation_color * alpha).astype(np.uint8) |
| 247 | fig, ax = plt.subplots() |
| 248 | ax.imshow(overlay) |
| 249 | patches = [plt.plot([], [], 's', color=np.array(color)/255, label=label)[0] for label, color in zip(label_names, colors)] |
| 250 | plt.legend(handles=patches, bbox_to_anchor=(0.5, -0.1), loc='upper center', ncol=4, fontsize='small') |
| 251 | plt.savefig(save_path, bbox_inches='tight') |
| 252 | plt.show() |
| 253 | |
| 254 | |
| 255 | def visualize_partition(coord, group_id, save_path): |
nothing calls this directly
no outgoing calls
no test coverage detected