How bit-ruiliu/mshnet works
This repository implements infrared small-target segmentation: a command-line training or testing run loads image/mask pairs, predicts masks with MSHNet, and evaluates detection quality. Training applies SLS loss to the final and (after warm-up) multi-scale predictions, updates model weights, and saves checkpoints; testing loads weights and reports metrics. The README also documents using SLS loss with other detectors, but no such detector implementations appear in the tree. Source-backed edges are cited; unsampled capabilities are not expanded into guessed components.
Main entry points: main.py Train/test runner · data.py IRSTD dataset · MSHNet.py MSHNet · loss.py SLS IoU loss
Components 17
Run orchestration
- Train/test runner
main.py - Weights and checkpoints
- Metrics output
- Adagrad optimizer
main.py
Detector network
- MSHNet
model/MSHNet.py - U-Net encoder
model/MSHNet.py - Middle layer
model/MSHNet.py - U-Net decoder
model/MSHNet.py - Multi-scale heads
model/MSHNet.py
Data and evaluation
- IRSTD dataset
utils/data.py - Paired transforms
utils/data.py - mIoU
utils/metric.py - PD/FA metric
utils/metric.py - ROC metric
utils/metric.py
Training objective
- SLS IoU loss
model/loss.py - Location penalty
model/loss.py
Other components
- Researcher
Connections 22
- Researcher → Train/test runner: starts run
README.md - Train/test runner → IRSTD dataset: builds loaders
main.py - IRSTD dataset → Paired transforms: transforms pairs
utils/data.py - Train/test runner → MSHNet: constructs and calls
main.py - MSHNet → U-Net encoder: encodes
model/MSHNet.py - MSHNet → Middle layer: processes bottleneck
model/MSHNet.py - MSHNet → U-Net decoder: decodes
model/MSHNet.py - MSHNet → Multi-scale heads: produces predictions
model/MSHNet.py - Train/test runner → SLS IoU loss: computes loss
main.py - SLS IoU loss → Location penalty: adds location penalty
model/loss.py - Train/test runner → Adagrad optimizer: updates weights
main.py - Train/test runner → mIoU: updates and reads
main.py - Train/test runner → PD/FA metric: updates and reads
main.py - Train/test runner → ROC metric: updates and reads
main.py - Train/test runner → Weights and checkpoints: loads and saves
main.py - Train/test runner → Metrics output: prints or logs
main.py - Multi-scale heads → SLS IoU loss: predictions scored
main.py - IRSTD dataset → MSHNet: supplies image batches
main.py - IRSTD dataset → mIoU: supplies labels
main.py - IRSTD dataset → PD/FA metric: supplies labels
main.py - IRSTD dataset → ROC metric: supplies labels
main.py - Metrics output → Researcher: reports results