3D CNNs for pollen classification

Overview

Telling the pollen of Parietaria and Urtica, two genera of the nettle family (Urticaceae), apart under the microscope is hard because the grains look very similar. A microscope captures each grain as a Z-stack: a series of images at different focal depths, each with a different part of the grain in focus. Earlier work classified these grains from 2D projections of the stack, which compress the depth information into a single image.

This thesis asks how well a network that sees the whole stack in 3D can classify the grains, and which parts of the stack it needs. It was my BSc thesis in Computer Science at Leiden University.

Data

The dataset holds about 600 Z-stacks, 200 per class, in three classes: Parietaria (P. judaica and P. officinalis), Urtica dioica/urens and Urtica membranacea. The stacks vary in both image size and number of frames.

Z-stack (variable size and depth)
  -> zero-pad or crop each frame to 224 x 224
  -> keep the middle 20 frames
  -> rescale to [-1, 1]
  -> training only: add augmented copies (flips, rotation)
  -> 3D CNN -> softmax over 3 classes

10 percent of the stacks are held out for testing and 10 percent of the remainder for validation, so the test set is never seen during training or model selection.

Method

Why 3D convolutions

A standard convolutional network slides 2D kernels over a single image. A 3D convolutional network uses kernels that also span the depth of the input, which is why such networks are used for volumetric images and video. Applied to a Z-stack, it can learn features that combine neighbouring focal planes, such as how a structure comes into and out of focus, rather than working from a single flattened view.

Preprocessing

Each frame is zero-padded or cropped to 224 × 224 pixels. The stack is then trimmed to its middle 20 frames, keeping the central part of the focal series, and pixel values are rescaled from [0, 1] to [−1, 1] to centre the inputs around zero. This gives every sample the same volumetric shape, which batched 3D convolution requires.

Augmentation

Augmentation adds copies rather than replacing originals. Each training stack is selected for an augmented copy with a probability set by an augmentation threshold, and a selected stack is randomly flipped horizontally and vertically and rotated by up to 100 degrees. Flips and rotations change the orientation of a grain but not its class, so they produce new samples without new labels. Keeping the originals means the model sees each grain at least once in its recorded orientation, and the threshold makes the amount of extra data a tunable hyperparameter.

Models

• Baseline 3D CNN. Two 3D convolution layers with 32 and 64 filters of size 3 × 3 × 3 and ReLU activations, with max pooling after the first, followed by dropout (rate 0.2), a fully connected hidden layer of 128 units and a three-class output. It serves as a simple reference point for what a shallow 3D network can learn from this data.
• ResNet3D. torchvision's r3d_18, an 18-layer residual network for video clips. Its stem convolution was replaced with one that accepts the 20-frame stack, and a final linear layer maps its 400 outputs to the three classes. The residual connections let the deeper network train stably, and for the main experiments it was trained from randomly initialised weights.

Training and evaluation

Both models were trained in PyTorch on Google Colab with cross-entropy loss and AdamW, whose decoupled weight decay acts as a regulariser. After every epoch the model was evaluated on the training and validation sets, and the checkpoint with the highest validation accuracy was kept for testing. The final configuration used a batch size of 16, a weight decay of 0.01 and 30 epochs.

Models were compared on accuracy and on weighted F1, which averages the per-class F1 scores in proportion to class size. F1 combines precision and recall, so it reflects whether the model separates all three classes rather than favouring one of them. Learning curves of loss, F1 and accuracy on the training and validation sets were used to diagnose overfitting.

Results

Baseline against ResNet3D

ResNet3D outperformed the baseline 3D CNN and was used for hyperparameter tuning and all remaining experiments.

Tuning

Of the three learning rates, 1e-4 gave the lowest test loss and the highest weighted F1 and accuracy.

• Augmentation. In a short three-epoch sweep over thresholds of 0, 0.25, 0.5 and 0.75, augmenting half of the training set gave the best test accuracy (0.79, against 0.70 to 0.71 for the other settings).
• Batch size. A batch size of 16 gave better test metrics than 32.

Whole stack against subsections

To test which focal planes the classifier needs, ResNet3D was also trained on subsections made of the middle frames only. The full 20-frame stack gave the best and most stable results, which indicates that the outer focal planes carry information that the central planes alone do not.

Transfer learning

Initialising ResNet3D from weights pretrained on video was tested as a way to compensate for the small dataset. It did not improve test accuracy, so the final model was trained from scratch.

Key takeaways

• The outer focal planes contribute information that the central planes lack, so classification benefits from the complete stack.
• A deeper 3D architecture (ResNet3D) outperformed a shallow baseline despite a dataset of only about 600 samples.
• Pretraining on video did not transfer to microscope stacks in this setting.

The project builds on Analysis of automatic image classification methods for Urticaceae pollen classification (Li et al., 2023).