Dual Input Fusion Classification

Overview

This repository contains a multi-input fusion deep learning pipeline for tumor grading from breast ultrasound. It combines B-mode ultrasound, the segmentation mask of the lesion, and strain elastography into a single classifier, and evaluates the result with stratified 5-fold cross-validation and fold ensembling.

Introduction

A B-mode ultrasound image alone carries the morphology of a lesion but little about its stiffness, while strain elastography carries stiffness but is noisy and hard to localize. A segmentation mask contributes neither texture nor stiffness, only where the lesion is. Each modality answers a different part of the grading question, so the pipeline fuses them at the input: the selected images are concatenated along the channel axis and projected back to three channels by a 1×1 convolution before entering a standard backbone. Any ImageNet-pretrained network can therefore be reused unchanged.

Two backbones are provided — a hierarchical vision transformer (HoVerTrans) and a lightweight CNN (ShuffleNetV2) — so the fusion strategy can be compared across model families.

Dataset

The dataset directory must follow this layout:

data/{dataset_name}/
├── train/
│   ├── img/           # B-mode ultrasound (loaded as grayscale)
│   ├── mask/          # segmentation mask (RGB)
│   ├── elastogram/    # strain elastography (RGB)
│   └── label.csv
└── test/
    ├── img/
    ├── mask/
    ├── elastogram/
    └── label.csv

The three subdirectories must share filenames — each row of label.csv is resolved against img/, mask/, and elastogram/ by the same name.

name,label
image001.png,0
image002.png,1
image003.png,2

Preprocessing. Images are resized to --img_size and normalized with mean 0.5 and standard deviation 0.5. During training the B-mode image additionally receives Gaussian noise and a 3×3 blur, each with 50% probability, and all three modalities receive the same horizontal flip with 50% probability so that they stay aligned. The B-mode image is read as single-channel and expanded to three channels before the model sees it.

Model architecture

Input fusion

--mode selects which modalities enter the network. A single-modality mode feeds its three channels straight to the backbone; a fusion mode concatenates the selected modalities on the channel axis and compresses them back to three channels with a 1×1 convolution, so the backbone always receives the same input shape.

Mode	Inputs	Channels
b	B-mode	3
mask	mask	3
se	elastography	3
b_mask	B-mode + mask	6 → 3
b_se	B-mode + elastography	6 → 3
mask_se	mask + elastography	6 → 3
b_mask_se	B-mode + mask + elastography	9 → 3

Running every mode through the same backbone makes the single-modality results directly comparable to the fusion results.

HoVerTrans

A hierarchical vision transformer with four stages. Each stage applies row-wise and column-wise attention through inner and outer transformer blocks, merges the pixel embeddings, and downsamples into the next stage. Stage depths, dimensions, and head counts are configurable from the command line.

CustomModel (ShuffleNetV2)

An ImageNet-pretrained ShuffleNetV2-x2.0 with its fully connected head replaced by one sized to --class_num, provided as an efficient CNN counterpart to HoVerTrans.

Training strategy

• Cross-validation: StratifiedKFold(n_splits=--fold, shuffle=True, random_state=42) over the training directory. Each fold trains its own model, and the held-out fold is used for validation with test-time transforms.
• Loss: cross-entropy with label smoothing 0.2.
• Optimizer: AdamW by default (SGD and Adam are also available) with a learning rate of 1e-4.
• Schedule: cosine annealing down to --min_lr, preceded by --warmup_epochs of linear warm-up. A step schedule is available through --scheduler step.
• Checkpointing: validation runs every --log_step epochs. The best checkpoint is saved whenever validation loss improves, but only after the first tenth of training has elapsed, which keeps an early noisy epoch from being selected. The final epoch is saved separately.
• All random seeds are fixed at 42 and cuDNN runs in deterministic mode.

Usage

pip install -r requirements.txt
pip install numpy pandas opencv-python scikit-learn matplotlib seaborn tqdm pillow tensorboard

requirements.txt pins only torch and torchvision; the packages on the second line are imported by the code but not listed there.

Training

CUDA_VISIBLE_DEVICES='0,1' python train.py \
    --data_path ./data/KO_20240805/KO/train \
    --class_num 5 \
    --model_name custom \
    --writer_comment KO_20240805/KO/b_se \
    --mode b_se

Testing

test.py loads the best checkpoint of every fold, evaluates each one on the test directory, and then combines them by both soft voting (mean of the predicted probabilities) and hard voting (majority of the predicted classes).

CUDA_VISIBLE_DEVICES='0,1' python test.py \
    --data_path ./data/KO_20240805/KO/test \
    --class_num 5 \
    --model_name custom \
    --writer_comment KO_20240805/KO/b_se \
    --mode b_se

Chain the two with && to train and evaluate in a single command.

Options

Option	Default	Description
--data_path	—	Dataset directory (train or test split)
--model_name	hovertrans	Backbone: hovertrans or custom
--mode	b	Input fusion mode
--class_num	2	Number of classes
--img_size	256	Input resolution
--batch_size	32	Batch size
--epochs	250	Training epochs
--lr	0.0001	Learning rate
--optimizer	AdamW	SGD, Adam, or AdamW
--scheduler	cosine	cosine or step
--fold	5	Number of cross-validation folds (training and testing alike)
--log_step	5	Validate every N epochs
--warmup_epochs	10	Linear warm-up length
--model_path	./weight	Checkpoint root
--writer_comment	GDPH&SYSUCC	Experiment identifier used in every output path

--writer_comment is what separates one experiment from another on disk, so give it a distinct value per run.

Metrics

Validation during training reports accuracy, sensitivity, specificity, precision, F1, and AUC (one-vs-rest for multiclass). Testing reports accuracy, macro F1, a per-class classification report, and row-normalized confusion matrices.

Outputs

weight/{model_name}/{writer_comment}/
├── model_info.txt              # the full argument set for this run
└── {fold}/
    ├── best_model.pth
    ├── last_epoch_model.pth
    ├── result_best.txt
    └── result_last_epoch.txt

logs/{model_name}/{writer_comment}_{fold}/    # TensorBoard

result/{model_name}/{writer_comment}/conf_matrix/
├── fold_1.png ... fold_N.png
├── ensemble_soft_voting.png
└── ensemble_hard_voting.png

Repository structure

├── config.py                     # argument parsing
├── dataset.py                    # dataset class and augmentation
├── model.py                      # HoVerTrans and CustomModel
├── train.py                      # cross-validated training
├── valid.py                      # validation loop and metrics
├── test.py                       # per-fold evaluation and ensembling
└── utils/
    ├── balance_augmentation.py   # class-balancing augmentation
    ├── label_*.py                # label.csv builders per dataset (BUSI, GDPH&SYSUCC, KO)
    ├── seg_preprocess/           # square padding, resize, grayscale, binary mask, [0,1] scaling
    └── strain_preprocess/        # DICOM to PNG, elastogram extraction, file organization, analysis

Notes

• Every visible GPU is wrapped in nn.DataParallel when more than one is present, so pin CUDA_VISIBLE_DEVICES to choose the devices a run uses. Checkpoints are always saved unwrapped, which keeps them loadable on a single GPU.
• --fold must match between training and testing: training writes one checkpoint directory per fold, and testing looks for exactly that many.
• data/, weight/, logs/, and result/ are git-ignored.

License

Dual Input Fusion Classification is released under the MIT License.