s3dis visualization code

[blank-track]

Hello, your work is excellent Can you provide the s3dis visualization code? thank you

[Gofinge]

Hi, here is our code for S3DIS visualization (Yixing's code), just for reference:

import open3d as o3d
from pathlib import Path
import torch
import numpy as np

from matplotlib import pyplot as plt
from scipy import stats

_label_to_color_uint8 = {
    0: [158, 218, 228],  # counter
    1: [151, 223, 137],  # floor
    2: [174, 198, 232],  # wall
    3: [255, 187, 120],  # bed
    4: [254, 127, 13],  # refrigerator
    5: [196, 176, 213],  # window
    6: [213, 39, 40],  # door
    7: [188, 189, 35],  # chair
    8: [255, 152, 151],  # table
    9: [140, 86, 74],  # sofa
    10: [196, 156, 147],  # bookshelf
    11: [148, 103, 188],  # picture
    12: [0, 0, 0],  # clutter
}

_label_to_color = dict([
    (label, (np.array(color_uint8).astype(np.float64) / 255.0).tolist())
    for label, color_uint8 in _label_to_color_uint8.items()
])

_name_to_color_uint8 = {
    "ceiling": [158, 218, 228],  # counter
    "floor": [151, 223, 137],  # floor
    "wall": [174, 198, 232],  # wall
    "beam": [255, 187, 120],  # bed
    "column": [254, 127, 13],  # refrigerator
    "window": [196, 176, 213],  # window
    "door": [213, 39, 40],  # door
    "chair": [188, 189, 35],  # chair
    "table": [255, 152, 151],  # table
    "sofa": [140, 86, 74],  # sofa
    "bookcase": [196, 156, 147],  # bookshelf
    "board": [148, 103, 188],  # picture
    "clutter": [0, 0, 0],  # clutter
}

_name_to_color = dict([(name, np.array(color_uint8).astype(np.float64) / 255.0)
                       for name, color_uint8 in _name_to_color_uint8.items()])

def load_real_data(pth_path):
    """
    Args:
        pth_path: Path to the .pth file.
    Returns:
        points: (N, 3), float64
        colors: (N, 3), float64, 0-1
        labels: (N, ), int64, {1, 2, ..., 36, 39, 255}.
    """
    # - points: (N, 3), float32           -> (N, 3), float64
    # - colors: (N, 3), float32, 0-255    -> (N, 3), float64, 0-1
    # - labels: (N, 1), float64, 0-19,255 -> (N,  ), int64, 0-19,255
    points, colors, labels = torch.load(pth_path)
    points = points.astype(np.float64)
    colors = colors.astype(np.float64) / 255.0
    assert len(points) == len(colors) == len(labels)

    labels = labels.astype(np.int64).squeeze()
    return points, colors, labels

def load_pred_labels(label_path):
    """
    Args:
        label_path: Path to the .txt file.
    Returns:
        labels: (N, ), int64, {1, 2, ..., 36, 39}.
    """
    def read_labels(label_path):
        labels = []
        with open(label_path, "r") as f:
            for line in f:
                labels.append(int(line.strip()))
        return np.array(labels)

    return np.array(read_labels(label_path))

def render_to_image(pcd, save_path):
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(pcd)
    vis.update_geometry(pcd)
    vis.poll_events()
    vis.update_renderer()
    vis.capture_screen_image(save_path)

def visualize_scene_by_path(scene_path, save_as_image=False):
    label_dir = Path("data/s3dis/test_epoch79")

    print(f"Visualizing {scene_path}")
    label_path = label_dir / f"{scene_path.stem}_pred.npy"

    # Load pcd and real labels.
    points, colors, real_labels = load_real_data(scene_path)

    # Visualize rgb colors
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    if save_as_image:
        render_to_image(pcd, f"image/{scene_path.stem}_rgb.png")
    else:
        o3d.visualization.draw_geometries([pcd], window_name="RGB colors")

    # Visualize real labels
    real_label_colors = np.array([_label_to_color[l] for l in real_labels])
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(real_label_colors)
    if save_as_image:
        render_to_image(pcd, f"image/{scene_path.stem}_real.png")
    else:
        o3d.visualization.draw_geometries([pcd], window_name="Real labels")

    # Load predicted labels
    pred_labels = np.load(label_path)
    pred_label_colors = np.array([_label_to_color[l] for l in pred_labels])
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(pred_label_colors)
    if save_as_image:
        render_to_image(pcd, f"image/{scene_path.stem}_pred.png")
    else:
        o3d.visualization.draw_geometries([pcd], window_name="Pred labels")

def visualize_scene_by_name(scene_name, save_as_image=False):
    data_root = Path("data") / "s3dis" / "Area_5"
    scene_paths = sorted(list(data_root.glob("*.pth")))

    found = False
    for scene_path in scene_paths:
        if scene_path.stem == scene_name:
            found = True
            visualize_scene_by_path(scene_path, save_as_image=save_as_image)
            break

    if not found:
        raise ValueError(f"Scene {scene_name} not found.")

if __name__ == "__main__":
    # Used in main text
    # hallway_10
    # lobby_1
    # office_27
    # office_30

    # Use in supplementary
    # visualize_scene_by_name("conferenceRoom_2")
    # visualize_scene_by_name("office_35")
    # visualize_scene_by_name("office_18")
    # visualize_scene_by_name("office_5")
    # visualize_scene_by_name("office_28")
    # visualize_scene_by_name("office_3")
    # visualize_scene_by_name("hallway_12")
    visualize_scene_by_name("office_4")

    # Visualize all scenes
    # data_root = Path("data") / "scannetv2" / "val"
    # scene_paths = sorted(list(data_root.glob("*.pth")))
    # scene_names = [p.stem for p in scene_paths]
    # for scene_name in scene_names:
    #     visualize_scene_by_name(scene_name, save_as_image=True)

[EricLina]

Nice work.Could you provide the code that analyzes the bad case?

[Gofinge]

Hi, maybe you can compute the mIoU of each Scene, then select a low mIoU scene and check visualization.

    data_root = Path("data") / "s3dis" / "Area_5"
    scene_paths = sorted(list(data_root.glob("*.pth")))
    label_dir = Path("data/s3dis/test_epoch79")
    out_path = Path("tools/s3dis_val_miou.csv")

    all_preds = []
    all_reals = []

    f = open(out_path, "w")
    f.write("scene_name,miou,K,entropy\n")

    for scene_path in scene_paths:
        label_path = label_dir / f"{scene_path.stem}_pred.npy"

        _, _, real_labels = load_real_data(scene_path)
        pred_labels = np.load(label_path)

        unique_labels = np.unique(real_labels)
        unique_K = len(unique_labels[unique_labels != 255])
        intersection, union, _ = intersection_and_union(
            output=np.array(pred_labels),
            target=np.array(real_labels),
            K=13,
            ignore_index=255)
        ious = intersection / (union + 1e-10)
        miou = np.sum(ious) / unique_K

        entropy = stats.entropy(real_labels)
        line = f"{scene_path.stem},{miou:.4f},{unique_K},{entropy:.4f}"
        print(line)
        f.write(line + "\n")

        all_preds.extend(pred_labels.tolist())
        all_reals.extend(real_labels.tolist())

    f.close()

    # Global
    intersection, union, _ = intersection_and_union(output=np.array(all_preds),
                                                    target=np.array(all_reals),
                                                    K=20,
                                                    ignore_index=255)
    ious = intersection / (union + 1e-10)
    miou = np.mean(ious)
    print(f"{label_path.stem} mIoU: {miou}")