Detection Datasets¶

Open in Colab · Download notebook

Detection datasets pair a scene with its boxes: a ragged \((K, 7)\) tensor and one class per box. Loaders that ship per-point instances instead of boxes can derive them.

Load a dataset¶

Provide the data directory containing the datasets, here we are using the environment variable TORCH_POINTCLOUD_DATA_DIR to point to this data directory. If you don't have the dataset locally in this directory, use download=True to download it.

Once downloaded in the raw directory, the dataset is automatically preprocessed in the processed directory.

data/
├── ModelNet40/
│   ├── raw/
│   │   ├── Area_1/
│   │   ├── Area_2/
│   │   └── ...
│   └── processed/
│       ├── Area_1/
│       ├── Area_2/
│       └── ...
└── ...

from torch_pointcloud.datasets import SunRGBD
from torch_pointcloud.config import DATA_DIR

dataset = SunRGBD(root=DATA_DIR, train=False, download=True)
print(f"Scenes: {len(dataset)}")
print({k: tuple(v.shape) for k, v in dataset[0].items()})

You can use force_process=True or force_download=True to redo either step, and num_workers to parallelize the processing.

What a sample holds¶

Some datasets add extra keys, depending on what they provide.

Batch the ragged boxes¶

Boxes cannot be concatenated like points without losing which scene they came from. Pass them as cat_keys, and collate writes a matching batch_box index.

from torch_pointcloud.utils.data import DataKeys, PointCloudDataLoader

dataloader = PointCloudDataLoader(
    dataset,
    batch_size=8,
    cat_keys=[DataKeys.BOX, DataKeys.LABEL],
)

data = next(iter(dataloader))
print(f"Box shape: {tuple(data['box'].shape)}")
print(f"Batch box shape: {tuple(data['batch_box'].shape)}")
print(f"Pos shape: {tuple(data['pos'].shape)}")

batch_box[i] names the scene of box i, as batch[j] names the scene of point j. Losses and metrics read both.

Driving datasets¶

Outdoor frames carry intensity instead of color, and their boxes are oriented. Some datasets require a manual download, as they require accepting terms of use first.

Boxes from instance labels¶

ScanNet ships per-point instance and segment rather than boxes. InstanceToBox turns them into detection targets: one axis-aligned box per instance, classed by its majority semantic label.

import numpy as np
import torch
from plyfile import PlyData

import torch_pointcloud.transforms as T

# Load the sample scene
ply = PlyData.read("sample_scene_labeled.ply")["vertex"]
pos = np.stack([ply["x"], ply["y"], ply["z"]], 1).astype("float32")
segment = np.asarray(ply["segment"]).astype("int64")
instance = np.asarray(ply["instance"]).astype("int64")

scene = {
    "pos": torch.from_numpy(pos),
    "segment": torch.from_numpy(segment),
    "instance": torch.from_numpy(instance),
}

# Derive one box per annotated instance
scene = T.InstanceToBox()(scene)
print(f"Box shape: {tuple(scene['box'].shape)}")
print(f"Label shape: {tuple(scene['label'].shape)}")

Every instance becomes a box here, walls and floor included. Map the stuff classes to the ignore_index with a Relabel first and their instances drop out of the box set. This is how the 18-class ScanNet detection targets are built:

from torch_pointcloud.datasets.scannet import SCANNET_DETECTION_LABELS

transform = T.Compose([
    T.Relabel(keys="segment", labels=SCANNET_DETECTION_LABELS, default=-1),
    T.InstanceToBox(),
])

transform