Skip to content

Point cloud segmentation with DGCNN

Author: K3-Node Team
Backend: Multi-Backend
Dataset: ShapeScenes
Description: Label every point of a point cloud.

View in Colab   GitHub source

Point cloud segmentation with DGCNN

Label every point of a point cloud. DGCNN (Wang et al., 2019) builds a new 30-nearest-neighbor graph in every layer, from the current point features, and learns from the differences between neighbors (DynamicEdgeConv). The features of all three layers are combined to classify each point.

Same model as PyG's examples/dgcnn_segmentation.py; on ShapeScenes instead of ShapeNet (not downloadable), with the mean IoU over the scene classes.

Install K3-Node, then choose a backend: "tensorflow", "torch" or "jax"

!pip install k3-node[examples]
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

Load the data

PyG's example segments ShapeNet airplanes into parts. ShapeScenes is a small stand-in: each scene holds three objects (cubes, spheres, cones or tori) and every point must be labeled with the kind of object it lies on (4 classes). The node features x are the surface normals. Training scenes are randomly jittered and rotated a little.

import keras
from keras import ops
from k3_node.datasets import ShapeScenes
from k3_node.loader import DataLoader
from k3_node.transforms import Compose, RandomJitter, RandomRotate
from k3_node.layers import DynamicEdgeConv
from k3_node.models import MLP

transform = Compose([RandomJitter(0.01), RandomRotate(15, axis=0), RandomRotate(15, axis=1), RandomRotate(15, axis=2)])
train_dataset = ShapeScenes("data/GeometricShapes", train=True, transform=transform)
test_dataset = ShapeScenes("data/GeometricShapes", train=False)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=10)
print(train_dataset[0])

Define the model

class Net(keras.Model):
    def __init__(self, out_channels, k=30, aggr="max"):
        super().__init__()
        self.conv1 = DynamicEdgeConv(MLP([2 * 6, 64, 64]), k, aggr)
        self.conv2 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
        self.conv3 = DynamicEdgeConv(MLP([2 * 64, 64, 64]), k, aggr)
        self.mlp = MLP([3 * 64, 1024, 256, 128, out_channels], dropout=0.5, norm=None)

    def call(self, data, training=False):
        x0 = ops.concatenate([data.x, data.pos], axis=-1)
        x1 = self.conv1(x0, data.batch, training=training)
        x2 = self.conv2(x1, data.batch, training=training)
        x3 = self.conv3(x2, data.batch, training=training)
        return self.mlp(ops.concatenate([x1, x2, x3], axis=1), training=training)


model = Net(train_dataset.num_classes)

Train

Every point is classified; besides the accuracy, the mean intersection-over-union (IoU) of the classes is reported, as in PyG. The learning rate is multiplied by 0.8 every 20 epochs.

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy", keras.metrics.MeanIoU(num_classes=4, sparse_y_pred=False, name="iou")],
)
model.fit(train_loader, validation_data=test_loader, epochs=30, callbacks=[keras.callbacks.LearningRateScheduler(lambda epoch: 0.001 * 0.8 ** (epoch // 20))],
          verbose=2)

Evaluate

loss, accuracy, iou = model.evaluate(test_loader, verbose=0)
print(f"Test accuracy: {accuracy:.4f}, IoU: {iou:.4f}")