Model & Training¶

Sparse 3D U-Net architecture, bridge point cloud dataset, and PyTorch Lightning training loop. Implements Phases D (model) and E (training) of the pipeline.

src.model ¶

Sparse ResNet U-Net Model for Bridge Point Cloud Classification

This module implements a sparse 3D U-Net architecture using spconv for efficient processing of sparse point cloud voxel grids. The model uses ResNet-style residual blocks in both encoder and decoder paths with skip connections.

Classes:

Name	Description
`ResidualBlock`	ResNet-style block with sparse convolutions
`SparseUNet`	U-Net architecture with encoder-decoder structure

ResidualBlock ¶

Bases: Module

ResNet-style residual block for sparse convolutions.

Uses SubMConv3d to maintain sparsity pattern while applying convolutions.

init ¶

__init__(in_channels: int, out_channels: int, norm_fn: Callable, indice_key: Optional[str] = None)

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of input channels	required
`out_channels`	`int`	Number of output channels	required
`norm_fn`	`Callable`	Normalization function (e.g., BatchNorm1d)	required
`indice_key`	`Optional[str]`	Key for spconv indice management	`None`

forward ¶

forward(x: SparseConvTensor) -> SparseConvTensor

Forward pass with residual connection.

Parameters:

Name	Type	Description	Default
`x`	`SparseConvTensor`	SparseConvTensor input	required

Returns:

Type	Description
`SparseConvTensor`	SparseConvTensor output

SparseUNet ¶

Bases: Module

Sparse 3D U-Net with ResNet-style blocks for bridge point cloud classification.

Architecture: - Encoder: 4 levels with downsampling (x1, x2, x4, x8) - Decoder: 3 levels with upsampling and skip connections - Output: Per-voxel classification into num_classes

init ¶

__init__(input_channels: int = 1, num_classes: int = 4, base_channels: int = 16)

Parameters:

Name	Type	Description	Default
`input_channels`	`int`	Number of input features (default: 1 for intensity)	`1`
`num_classes`	`int`	Number of output classes (default: 4) - 0: Background/Unclassified - 1: Ground/Water - 2: Bridge Deck - 3: Obstacles/High Noise	`4`
`base_channels`	`int`	Base number of channels (default: 16)	`16`

freeze_encoder ¶

freeze_encoder() -> None

Freeze all encoder layers. Only decoder and classifier remain trainable.

forward ¶

forward(x: SparseConvTensor) -> Tensor

Forward pass through the U-Net.

Parameters:

Name	Type	Description	Default
`x`	`SparseConvTensor`	spconv.SparseConvTensor with features and indices	required

Returns:

Type	Description
`Tensor`	torch.Tensor: (N, num_classes) logits for each voxel

src.dataset ¶

Bridge point cloud dataset with on-the-fly voxelization.

Provides BridgeDataset (PyTorch Dataset) and sparse_collate_fn for batching voxelized point clouds into sparse tensor format.

BridgeDataset ¶

Bases: Dataset

Dataset for bridge point cloud classification with voxelization.

Handles HUC-organized directory structure and properly aggregates points within voxels using majority vote for labels and averaging for features.

init ¶

__init__(data_dir: str, voxel_size: float = 0.1, augment: bool = False, augment_extra: bool = False, max_voxels: Optional[int] = None)

Parameters:

Name	Type	Description	Default
`data_dir`	`str`	Path to directory containing .npy files (can be HUC-organized).	required
`voxel_size`	`float`	Voxel size in meters (e.g., 0.1 for 10cm).	`0.1`
`augment`	`bool`	Whether to apply random Z-rotation and jitter.	`False`
`augment_extra`	`bool`	Extra augmentation (XY-flip, scaling, intensity jitter, point dropout). Requires augment=True.	`False`
`max_voxels`	`Optional[int]`	Maximum voxels per sample; randomly subsample if exceeded (default: None = no limit).	`None`

getitem ¶

__getitem__(idx: int) -> Tuple[ndarray, ndarray, ndarray]

Load and voxelize a single bridge sample.

Returns:

Type	Description
`ndarray`	Tuple of (discrete_coords, features, labels)
`ndarray`	discrete_coords: (N_voxels, 3) integer voxel coordinates
`ndarray`	features: (N_voxels, 1) averaged intensity per voxel
`Tuple[ndarray, ndarray, ndarray]`	labels: (N_voxels,) majority-vote labels per voxel

sparse_collate_fn ¶

sparse_collate_fn(batch: List[Tuple[ndarray, ndarray, ndarray]]) -> Dict[str, Tensor]

Custom collate function to create a batch for sparse tensor format.

Sparse tensors require coordinate format: [Batch_ID, X, Y, Z]

Parameters:

Name	Type	Description	Default
`batch`	`List[Tuple[ndarray, ndarray, ndarray]]`	List of (coords, features, labels) tuples	required

Returns:

Type	Description
`Dict[str, Tensor]`	Dictionary with keys:
`Dict[str, Tensor]`	coordinates: (N_total, 4) tensor [batch_id, x, y, z]
`Dict[str, Tensor]`	features: (N_total, 1) tensor
`Dict[str, Tensor]`	labels: (N_total,) tensor
`Dict[str, Tensor]`	sample_voxel_counts: list of int

src.train ¶

Bridge Classification Model Training - Data Loader with Voxelization.

This module provides a data loader for bridge point cloud classification with:

Proper voxelization with feature aggregation and majority-vote labeling
Support for HUC-organized directory structure
Visualization tools to verify voxelization correctness
PyTorch Lightning training integration

Data: use --train-dir for training data and optionally --val-dir for validation. If --val-dir is not provided, the script uses --val-split to randomly split the training directory into train/validation (val-split=0 means no validation). Testing (gold/human-labeled) data is not used during training; use it later for final evaluation after model selection.

Usage

# Test data loader (default: --train-dir ./data/ml-data/training)
python src/train.py --train-dir ./data/ml-data/training

# Visualize voxelization
python src/train.py --visualize --sample-idx 0

# Train with explicit validation directory
python src/train.py --train --train-dir ./data/ml-data/training         --val-dir ./data/ml-data/validation --epochs 50

# Train with val-split when val-dir not provided
python src/train.py --train --train-dir ./data/ml-data/training         --val-split 0.2 --epochs 50

DiceLoss ¶

Bases: Module

Per-class Dice loss averaged across classes. Directly optimizes overlap (IoU proxy).

BridgeLightningModule ¶

Bases: LightningModule

PyTorch Lightning module for bridge classification training.

Wraps SparseUNet with weighted cross-entropy (optionally + Dice) loss, deck-specific metrics (IoU, precision, recall), and ReduceLROnPlateau scheduling.

init ¶

__init__(input_channels: int = 1, num_classes: int = NUM_CLASSES, base_channels: int = 16, learning_rate: float = 0.001, weight_decay: float = 0.01, class_weights: list | None = None, monitor: str = 'val_loss', monitor_mode: str = 'min', use_dice_loss: bool = False)

Parameters:

Name	Type	Description	Default
`input_channels`	`int`	Number of input features (default: 1).	`1`
`num_classes`	`int`	Number of output classes (default: 4).	`NUM_CLASSES`
`base_channels`	`int`	Base number of channels for the U-Net encoder (default: 16).	`16`
`learning_rate`	`float`	Learning rate for AdamW optimizer (default: 0.001).	`0.001`
`weight_decay`	`float`	Weight decay for AdamW optimizer (default: 0.01).	`0.01`
`class_weights`	`list \| None`	Per-class weights for loss function. If None, uses defaults from calculate_weights.py.	`None`
`monitor`	`str`	Metric name for ReduceLROnPlateau and checkpointing (default: val_loss).	`'val_loss'`
`monitor_mode`	`str`	'min' or 'max' for ReduceLROnPlateau (default: min).	`'min'`
`use_dice_loss`	`bool`	If True, use combined 0.5CE + 0.5Dice loss (default: False).	`False`

forward ¶

forward(x: SparseConvTensor) -> Tensor

Forward pass.

training_step ¶

training_step(batch: dict, batch_idx: int) -> Optional[Tensor]

Training step.

validation_step ¶

validation_step(batch: dict, batch_idx: int) -> Optional[Tensor]

Validation step.

configure_optimizers ¶

configure_optimizers() -> dict

Configure optimizer.

BridgeDataModule ¶

Bases: LightningDataModule

PyTorch Lightning data module for bridge dataset.

Uses train_dir and optionally val_dir, or val_split on train_dir when val_dir not set. Handles voxelization config, augmentation flags, and DataLoader setup.

init ¶

__init__(train_dir: str, val_dir: Optional[str] = None, voxel_size: float = 0.1, batch_size: int = 4, num_workers: int = 4, augment: bool = True, augment_extra: bool = False, val_split: float = 0.0, max_voxels: Optional[int] = None)

Parameters:

Name	Type	Description	Default
`train_dir`	`str`	Path to directory containing training .npy files	required
`val_dir`	`Optional[str]`	Path to directory containing validation .npy files; if None, use val_split on train_dir	`None`
`voxel_size`	`float`	Voxel size in meters (default: 0.1)	`0.1`
`batch_size`	`int`	Batch size (default: 4)	`4`
`num_workers`	`int`	Number of data loader workers (default: 4)	`4`
`augment`	`bool`	Whether to apply augmentation (default: True)	`True`
`augment_extra`	`bool`	Whether to apply extra augmentation (default: False)	`False`
`val_split`	`float`	Validation split ratio when val_dir is not set (default: 0.0, no validation)	`0.0`
`max_voxels`	`Optional[int]`	Maximum voxels per sample (default: None = no limit)	`None`

setup ¶

setup(stage=None)

Setup datasets from train_dir and val_dir, or val_split on train_dir.

train_dataloader ¶

train_dataloader()

Create training dataloader.

val_dataloader ¶

val_dataloader()

Create validation dataloader.

save_network_graph ¶

save_network_graph(model: Module, save_dir: str, filename: str = 'network_architecture') -> None

Traces the model on GPU (required for spconv) and saves the graph as PNG.

visualize_voxelization ¶

visualize_voxelization(data_dir: str, sample_idx: int = 0, voxel_size: float = 0.1) -> None

Visualize original vs voxelized point cloud for a sample bridge.

Parameters:

Name	Type	Description	Default
`data_dir`	`str`	Path to data directory	required
`sample_idx`	`int`	Index of sample to visualize	`0`
`voxel_size`	`float`	Voxel size in meters	`0.1`

main ¶

main() -> None

Main entry point with command-line argument parsing.

train.py CLI Arguments¶

Argument	Default	Description
`--data-dir`	`./data/ml-data`	Data directory for test loader and visualize when not using --train-dir
`--train-dir`	`./data/ml-data/training`	Training data directory
`--val-dir`	None	Validation directory; if unset, uses `--val-split`
`--val-split`	0.0	Fraction of training data to use as validation (0 = none)
`--voxel-size`	0.1	Voxel size in meters
`--max-voxels`	None	Max voxels per sample; subsampled if exceeded (OOM prevention)
`--batch-size`	16	Batch size
`--augment`	False	Enable random Z-rotation + jitter
`--augment-extra`	False	Extra augmentation: XY-flip, scaling, intensity jitter, point dropout. Requires `--augment`
`--dice-loss`	False	Use combined Dice + CrossEntropy loss (0.5CE + 0.5Dice)
`--train`	False	Enable training mode
`--epochs`	50	Training epochs
`--learning-rate`	0.001	AdamW learning rate
`--weight-decay`	0.01	AdamW weight decay
`--base-channels`	16	U-Net base channel count
`--num-workers`	4	DataLoader workers
`--exp-name`	`bridge_classify_base`	Experiment name for logs/checkpoints
`--experiments-dir`	`./experiments`	Base directory for experiments
`--class-weights`	None	Path to `class_weights.json` from `calculate_weights.py`
`--gpus`	auto	Number of GPUs (None = auto-detect). GPU required.
`--early-stopping`	False	Stop when monitored metric stops improving
`--early-stopping-patience`	10	Epochs to wait before early stopping
`--monitor`	`val_deck_iou`	Metric for checkpointing + early stopping
`--accumulate-grad-batches`	1	Gradient accumulation steps
`--ckpt-path`	None	Checkpoint path to resume training from
`--finetune`	None	Checkpoint for fine-tuning (weights only, fresh optimizer/epoch). Mutually exclusive with `--ckpt-path`
`--freeze-encoder`	False	Freeze encoder; only decoder and classifier are trained
`--save-top-k`	5	Number of best checkpoints to keep
`--visualize`	False	Visualize voxelization for a sample
`--sample-idx`	0	Sample index to visualize

train.py Usage Examples¶

# Basic training
python src/train.py --train \
    --train-dir ./data/ml-data/training \
    --val-dir ./data/ml-data/validation \
    --class-weights ./data/ml-data/class_weights.json \
    --exp-name my-experiment \
    --epochs 50 --batch-size 16 --augment

# Resume from checkpoint
python src/train.py --train \
    --train-dir ./data/ml-data/training \
    --val-dir ./data/ml-data/validation \
    --class-weights ./data/ml-data/class_weights.json \
    --ckpt-path ./experiments/my-experiment/version_0/checkpoints/last.ckpt \
    --exp-name my-experiment-resumed

# Fine-tune from a pretrained model (fresh optimizer, new experiment)
python src/train.py --train \
    --train-dir ./data/ml-data/gold-data-normalized \
    --val-dir ./data/ml-data/gold-data-normalized-val \
    --finetune ./experiments/base-model/version_0/checkpoints/best.ckpt \
    --exp-name fine-tuned-gold \
    --epochs 30 --learning-rate 0.0005 --freeze-encoder

Model & Training¶

src.model ¶

ResidualBlock ¶

__init__ ¶

forward ¶

SparseUNet ¶

__init__ ¶

freeze_encoder ¶

forward ¶

src.dataset ¶

BridgeDataset ¶

__init__ ¶

__getitem__ ¶

sparse_collate_fn ¶

src.train ¶

DiceLoss ¶

BridgeLightningModule ¶

__init__ ¶

forward ¶

training_step ¶

validation_step ¶

configure_optimizers ¶

BridgeDataModule ¶

__init__ ¶

setup ¶

train_dataloader ¶

val_dataloader ¶

save_network_graph ¶

visualize_voxelization ¶

main ¶

train.py CLI Arguments¶

train.py Usage Examples¶

init ¶

init ¶

init ¶

getitem ¶

init ¶

init ¶