Pipe Material Classification

A comprehensive image classification project for classifying different types of pipe materials using deep learning.

Project Structure

.
├── conf/                          # Hydra configuration files
│   ├── config.yaml               # Main configuration
│   └── model/                    # Model configurations
│       ├── segformer.yaml
│       └── unet.yaml
├── data/                         # Data directories
│   ├── dataset/                 # Training/validation/test datasets
│   │   ├── acciaio/
│   │   ├── acciaio_rivestito/
│   │   ├── ghisa/
│   │   ├── polietilene/
│   │   └── PVC/
│   └── modules/                 # Saved model checkpoints and modules
├── outputs/                      # Training outputs
│   ├── checkpoints/             # Model checkpoints
│   └── predictions/             # Prediction results
├── src/pipe_material/
│   ├── data/                    # Data loading modules
│   │   ├── dataset/
│   │   │   ├── dataset.py             # Single-label Dataset class
│   │   │   └── multi_dataset.py       # Multi-label folder-based Dataset classes
│   │   └── datamodule/
│   │       ├── datamodule.py          # Single-label PyTorch Lightning DataModule
│   │       └── multi_datamodule.py    # Multi-label PyTorch Lightning DataModule
│   ├── models/                  # Model implementations
│   │   ├── base.py             # Base classifier class
│   │   ├── segformer.py        # SegFormer model
│   │   ├── unet.py             # UNet model
│   │   ├── resnet.py           # ResNet model
│   │   └── clip_vit.py         # CLIP ViT model
│   ├── training/                # Training scripts
│   │   ├── trainer.py          # Single-label training script
│   │   └── multi_trainer.py    # Multi-label (folder-level) training script
│   ├── cli.py                  # CLI interface
│   └── core.py                 # Core utilities
├── notebooks/                   # Jupyter notebooks
│   └── example.ipynb
├── tools/                       # Additional tools
└── README.md                    # This file

Installation

Using uv (recommended)

uv sync --all-groups

Using pip

pip install -r requirements.txt

Dataset Preparation

Organize your data in the following structure:

data/dataset/
├── class_1/
│   ├── image_1.jpg
│   ├── image_2.jpg
│   └── ...
├── class_2/
│   ├── image_1.jpg
│   ├── image_2.jpg
│   └── ...
└── class_n/
    └── ...

Example with pipe materials:

data/dataset/
├── acciaio/
│   ├── pipe_001.jpg
│   ├── pipe_002.jpg
│   └── ...
├── acciaio_rivestito/
│   └── ...
├── ghisa/
│   └── ...
├── polietilene/
│   └── ...
└── PVC/
    └── ...

Multi-Label Classification (Folder-Level Labels)

In addition to the single-label, one-class-per-directory setup above, the project supports a folder-level, multi-label workflow for cases where a single sample (e.g. a pipe inspection) is represented by a folder of images that can belong to more than one material class at once.

This workflow is implemented in multi_dataset.py, multi_datamodule.py, and multi_trainer.py.

Expected data layout

Data is organized as one numbered folder per sample, each containing all images for that sample, plus a single CSV mapping folder IDs to a list of class labels:

data/
├── train/
│   ├── 1/
│   │   ├── img_a.jpg
│   │   └── img_b.jpg
│   ├── 2/
│   │   └── img_c.jpg
│   └── ...
├── val/
│   └── ...
├── test/
│   └── ...
└── labels.csv

labels.csv uses a ;-delimited id,labels format, where labels is a list of class ids (any separator inside the field is tolerated, e.g. "0;2", "0,2", "[0, 2]" all parse to [0, 2]):

id;labels
1;0;2
2;1

Each class id present for a folder is converted into a multi-hot label vector of length num_classes (inferred automatically from the CSV as max(label_id) + 1 if not set explicitly).

Two related components

• ImageFolderMultiLabelDataset / ImageClassificationMultiLabelDataModule (multi_dataset.py, multi_datamodule.py) — a per-image, validation-only dataset/datamodule. Each item is a single image plus the id of the folder it belongs to (no train_dataloader/test_dataloader, only val_dataloader). It's meant for a workflow where a model scores each image independently, predictions are grouped back by folder_id, and compared as a set against the folder's ground truth via get_ground_truth / get_ground_truth_multi_hot.
• FolderMultiLabelDataset / FolderMultiLabelDataModule (in multi_trainer.py) — a folder-level (bag-of-images) dataset used for actual training. Each item returns all images in a folder as a single stacked tensor together with the folder's multi-hot target. A custom collate function keeps folders as a list (since folders can contain a variable number of images) rather than stacking them into a fixed-size batch tensor.

Training a multi-label model

multi_trainer.py wraps any of the project's per-image classifiers (segformer, unet, clip_vit, resnet) inside MultiLabelFolderClassifier, a pytorch_lightning.LightningModule that:

1. Runs the chosen classifier on every image in a folder to get per-image logits.
2. Pools those logits into a single folder-level logit vector via a log-sum-exp pooling (a smooth approximation of max-pooling across images).
3. Optimizes a multi-label BCEWithLogitsLoss against the folder's multi-hot ground truth.
4. Reports set-based precision, recall, F1, and exact-match metrics per folder (treating predictions above a probability threshold, default 0.5, as the predicted label set).

Run training with:

uv run python -m pipe_material.training.multi_trainer

or directly:

uv run python src/pipe_material/training/multi_trainer.py

This uses Hydra with conf/multi_train.yaml as the default config. Override parameters the same way as for the single-label trainer, e.g.:

uv run python src/pipe_material/training/multi_trainer.py model=unet training.batch_size=4

Multi-label configuration (conf/multi_train.yaml)

defaults:
  - model: resnet

training:
  batch_size: 8
  num_epochs: 100
  learning_rate: 0.001
  weight_decay: 0.0001
  num_workers: 4
  pin_memory: true
  seed: 42

data:
  train_dir: /path/to/data/train
  val_dir: /path/to/data/val
  test_dir: /path/to/data/val
  train_labels_csv: /path/to/data/labels.csv
  val_labels_csv: /path/to/data/labels.csv
  test_labels_csv: /path/to/data/labels.csv
  image_size: 224
  augmentation: true

logger:
  project_name: "pipe-material-classification"
  log_frequency: 10

checkpoint:
  monitor: "val_loss"
  mode: "min"
  save_dir: ./outputs/checkpoints

device:
  accelerator: "gpu"
  devices: 1
  precision: "16-mixed"

Training logs to Weights & Biases via WandbLogger, checkpoints the top-3 models plus the last epoch (ModelCheckpoint), applies early stopping (patience 20) and learning-rate monitoring, and automatically runs a test-set evaluation once training completes.

Configuration

The project uses Hydra for configuration management. Main configuration is in conf/config.yaml.

Key Configuration Parameters

training:
  batch_size: 32 # Batch size for training
  num_epochs: 100 # Number of training epochs
  learning_rate: 0.001 # Learning rate
  num_workers: 4 # Number of data loading workers
  seed: 42 # Random seed

data:
  data_dir: ./data/dataset
  train_split: 0.8 # 80% training
  val_split: 0.1 # 10% validation
  test_split: 0.1 # 10% testing
  image_size: 224 # Input image size
  augmentation: true # Enable data augmentation

model:
  _target_: pipe_material.models.segformer.SegFormerClassifier
  num_classes: 5
  pretrained: true

device:
  accelerator: "gpu" # 'gpu' or 'cpu'
  devices: 1 # Number of GPUs
  precision: "16-mixed" # Mixed precision training

Training

Basic Training

Train with default configuration:

uv run -m pipe_material train

Training with Custom Configuration

Override specific parameters:

uv run -m pipe_material train training.batch_size=64 training.learning_rate=0.0005

Switch Model Architecture

Train with UNet instead of SegFormer:

uv run -m pipe_material train --config-path conf --config-name config model=unet

Outputs

Training outputs are saved to outputs/:

outputs/
├── checkpoints/
│   ├── best-epoch00-val_loss0.25.ckpt
│   ├── best-epoch05-val_loss0.18.ckpt
│   ├── last.ckpt
│   └── ...
├── logs/
│   └── (wandb logs)
└── predictions/
    └── (inference results)