Pivot Robotics Computer Vision Challenge

Implementation of the 3D Student-Teacher (3D-ST) method for anomaly detection of 3D data as described in this paper.

Synthetic Dataset Generation

A synthetic dataset is generated using the method described in page 9 of the paper consisting of a training and validation set of 500 and 25 point clouds using the ModelNet10 dataset. The number of points sampled per point cloud is set to 16,000 instead of the 64,000 as in the original paper to save time.

python -m modules.data.create_datasets

Teacher and Decoder Pretraining

Because of the smaller point_cloud size, k is reduced to 8 from 32 and m is reduced to 128. The feature vectors are 64-dimensional and there are 4 residual network blocks. The models were trained for 250 epochs, where at each time step, 16 random feature vectors are sampled from the teacher's output and fed to the decoder. The total chamfer distance between the decoder's output clouds and the receptive field clouds is used as the loss for both models.

python -m modules.pretrain_teacher

Student Training

We train the student model using the same architecture as the teacher. The loss is calculated via the average L2 distance between the student's and teacher's feature vectors. Since the student is only being trained on anomaly-free data, the idea is that it will only learn to predict features that are anomaly free, and the error from the teacher can be measured as a form of heatmap for defects.

python -m modules.train_student

Inference

python -m modules.create_heatmap

Examples

Crack anomalies

Hole defect

Anomaly-free

Teacher training loss

Student network training loss