Part Segmentation¶
Open in Colab · Download notebook
Part segmentation labels every point of a single object with one of its parts: a chair's legs, back and seat. Models read a packed batch and the object category, and return per-point logits.
ShapeNetPart, the standard benchmark, spreads 50 part ids over 16 categories.
Run a pretrained checkpoint¶
The example reads one airplane of the ShapeNetPart test set. The dataset has no automatic download: get it from shapenet.org and extract it under data/ShapeNetPart/raw/.
Pass return_info=True to get the checkpoint's transform: it samples the point cloud, carries normal and segment along, and one-hots the category.
import torch
import torch_pointcloud as tp
from torch_pointcloud.datasets import ShapeNetPart
from torch_pointcloud.utils.data import collate
# Load the pretrained model
model, info = tp.create_model(
"pointnext-sm.shapenetpart.openpoints",
task="segmentation",
pretrained=True,
return_info=True,
)
model = model.eval()
# Get associated transform
transform = info["transform"]
# Load one airplane of the ShapeNetPart test set
dataset = ShapeNetPart(root="data", split="test", categories="Airplane", transform=transform)
sample = dataset[0]
# Collate the sample into a batch
data = collate([sample])
print(f"Data keys: {data.keys()}")
# Inference pass, the category enters as a fourth argument
with torch.no_grad():
logits = model(data.get("x"), data["pos"], data["batch"], data["category"])
# Get predictions inside the category's own part ids
part_ids = ShapeNetPart.seg_ids["Airplane"]
preds = logits[:, part_ids].argmax(dim=-1)
print(f"Logits shape: {tuple(logits.shape)}")
print(f"Predicted parts: {preds.unique(return_counts=True)}")
# Data keys: dict_keys(['pos', 'normal', 'segment', 'category', 'origin_pos', 'origin_segment', 'index', 'height', 'x', 'batch'])
# Logits shape: (2048, 50)
# Predicted parts: (tensor([0, 1, 2, 3]), tensor([1044, 542, 219, 243]))
The head scores all 50 part ids at once. The reporting protocol argmaxes inside the four ids the airplane owns, so preds counts its body, wing, tail and engine points rather than indexing the global 50.
Pass a placeholder segment
The transform subsamples the part labels alongside the points, so it expects a segment key even at inference. Any tensor of the right length does.
Inputs and outputs¶
| Argument | Shape | Description |
|---|---|---|
x |
\((N, C)\) or None |
Per-point features, usually normals and height |
pos |
\((N, 3)\) | Coordinates, all objects in the batch concatenated |
batch |
\((N,)\) | Index tensor associating each point with its object |
category |
\((B, 16)\) | One-hot object category, one row per object |
| returns | \((N, 50)\) | Per-point logits over the 50 part ids |
category is per object, not per point: collate stacks it to \((B, 16)\) while pos is concatenated to \((N, 3)\).
Category ids and part ids¶
from torch_pointcloud.datasets import ShapeNetPart
print(len(ShapeNetPart.category_ids), list(ShapeNetPart.category_ids)[:4])
print(ShapeNetPart.seg_ids["Chair"], ShapeNetPart.seg_ids["Table"])
# 16 ['Airplane', 'Bag', 'Cap', 'Car']
# [12, 13, 14, 15] [47, 48, 49]
category_ids is ordered, so list(ShapeNetPart.category_ids).index(name) is the integer the one-hot encodes, and seg_ids[name] is the slice of the 50 outputs that category owns.
Evaluate on a dataset¶
The benchmark metric is instance mIoU: per object, the mean IoU over its own parts, averaged over objects. A part absent from both the prediction and the target counts as 1.0. Class mIoU averages per category first.
The scoring helpers live in torch_pointcloud.metrics.
import torch
from tqdm.auto import tqdm
import torch_pointcloud as tp
from torch_pointcloud.datasets import ShapeNetPart
from torch_pointcloud.utils.data import PointCloudDataLoader
from torch_pointcloud.metrics import part_intersection_over_union, part_mean_intersection_over_union
from torch_pointcloud.config import DATA_DIR
model, info = tp.create_model(
"pointnext-sm.shapenetpart.openpoints",
task="segmentation",
pretrained=True,
return_info=True,
)
model = model.cuda().eval()
dataset = ShapeNetPart(root=DATA_DIR, split="test", transform=info["transform"])
dataloader = PointCloudDataLoader(dataset, batch_size=16, num_workers=6)
part_ids = list(ShapeNetPart.seg_ids.values())
ious, categories = [], []
with torch.no_grad():
for data in tqdm(dataloader, desc="Evaluating"):
category = data["category"].cuda()
batch = data["batch"].cuda()
logits = model(data["x"].cuda(), data["pos"].cuda(), batch, category)
preds = logits.argmax(dim=-1)
index = category.argmax(dim=-1)
ious.append(part_intersection_over_union(preds, data["segment"].cuda(), part_ids, index, batch))
categories.append(index)
ious, categories = torch.cat(ious), torch.cat(categories)
instance_miou = part_mean_intersection_over_union(ious, categories)
class_miou = part_mean_intersection_over_union(ious, categories, average="macro")
print(f"instance mIoU {instance_miou:.4f} | class mIoU {class_miou:.4f}")
# instance mIoU 0.8587 | class mIoU 0.8348
Train from scratch¶
The loop below trains pointnext-sm.shapenetpart.openpoints from scratch on the ShapeNetPart train split, for a single demo epoch. The transform one-hot encodes category, which the model takes alongside the points, and the loss runs over the 50 part ids.
from tqdm.auto import tqdm
from torch.nn import functional as F
import torch
import torch_pointcloud as tp
import torch_pointcloud.transforms as T
from torch_pointcloud.datasets import ShapeNetPart
from torch_pointcloud.utils.data import PointCloudDataLoader
from torch_pointcloud.config import DATA_DIR
device = "cuda"
# Setup the dataset and dataloader
train_dataset = ShapeNetPart(
DATA_DIR,
split="train",
transform=T.Compose([
T.FarthestPointSample(keys=["pos", "normal", "segment"], pos_key="pos", num_samples=2048),
T.Rescale(keys="pos", method="centroid"),
T.Cat(keys=["pos", "normal"], dst_key="x"),
T.OneHot(keys="category", num_classes=16),
]),
)
train_dataloader = PointCloudDataLoader(
train_dataset,
batch_size=16,
shuffle=True,
num_workers=4,
)
# Create the desired model and optimizer
model = tp.create_model(
"pointnext-sm.shapenetpart.openpoints",
task="segmentation",
in_channels=6,
num_classes=50,
num_categories=16,
).to(device)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
weight_decay=1e-4,
)
# Training loop
num_epochs = 1 # a demo pass, raise it to actually train
model.train()
for epoch in range(num_epochs):
total_loss = 0.0
pbar = tqdm(enumerate(train_dataloader), total=len(train_dataloader), desc=f"Epoch {epoch}")
for i, data in pbar:
x = data["x"].to(device)
pos = data["pos"].to(device)
target = data["segment"].to(device)
batch = data["batch"].to(device)
category = data["category"].to(device)
optimizer.zero_grad()
logits = model(x, pos, batch, category)
loss = F.cross_entropy(logits, target)
loss.backward()
optimizer.step()
total_loss += loss.item()
if (i + 1) % 10 == 0:
loss_step = loss.item()
metrics = {"train/loss_step": f"{loss_step:.3f}"}
pbar.set_postfix(metrics)
loss_epoch = total_loss / len(train_dataloader)
print(f"Loss epoch {epoch}: {loss_epoch:.3f}")