make_oracle.py

"""Build oracle (ground-truth) 3D instance masks for ScanNet200.

Converts ScanNet200 ground-truth instance `.txt` files into `.pth` mask arrays
used as "oracle" 3D proposals (see paper Table II). Instance id 0 (unlabeled)
is remapped to -1; remaining ids are compacted to a 0-based range.

Usage:
    python make_oracle.py \
        --gt-dir  data/scannetv2/scannet200_gt/instance_gt/validation \
        --out-dir output/scannet200/oracle/inst_seg_pcd
"""
import argparse
import glob
import os

import numpy as np
import torch


def transform_data(data):
    """Remap instance ids: 0 -> -1, others -> compact 0-based ids."""
    unique_vals = np.unique(data)
    mapping = {0: -1}
    next_val = 0
    for val in unique_vals:
        if val != 0:
            mapping[val] = next_val
            next_val += 1
    return np.vectorize(mapping.get)(data)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--gt-dir",
        default="data/scannetv2/scannet200_gt/instance_gt/validation",
        help="Directory of ScanNet200 GT instance `scene*.txt` files.",
    )
    parser.add_argument(
        "--out-dir",
        default="output/scannet200/oracle/inst_seg_pcd",
        help="Directory to write oracle `.pth` mask arrays.",
    )
    args = parser.parse_args()

    scene_list = glob.glob(os.path.join(args.gt_dir, "scene*.txt"))
    if not scene_list:
        raise FileNotFoundError(f"No `scene*.txt` files found in {args.gt_dir}")
    os.makedirs(args.out_dir, exist_ok=True)

    for scene_path in scene_list:
        gt_data = np.loadtxt(scene_path, delimiter=None, dtype=int)
        transformed_data = transform_data(gt_data)
        scene_name = os.path.splitext(os.path.basename(scene_path))[0]
        save_path = os.path.join(args.out_dir, f"{scene_name}.pth")
        torch.save(transformed_data, save_path)
        print(f"Processed and saved: {save_path}")


if __name__ == "__main__":
    main()