# bit-ruiliu/mshnet architecture

> Architecture diagram of the GitHub repository bit-ruiliu/mshnet, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-30). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/bit-ruiliu/mshnet
- Repository: https://github.com/bit-ruiliu/mshnet

## Overview

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`](https://github.com/bit-ruiliu/mshnet/blob/main/main.py) (Train/test runner), [`utils/data.py`](https://github.com/bit-ruiliu/mshnet/blob/main/utils/data.py) (IRSTD dataset), [`model/MSHNet.py`](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py) (MSHNet), [`model/loss.py`](https://github.com/bit-ruiliu/mshnet/blob/main/model/loss.py) (SLS IoU loss)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_run["Run orchestration"]
  node_runner["Train/test runner<br/>[main.py]"]
  node_weights["Weights and checkpoints"]
  node_report["Metrics output"]
  node_optimizer["Adagrad optimizer<br/>[main.py]"]
end

subgraph group_network["Detector network"]
  node_mshnet["MSHNet<br/>[MSHNet.py]"]
  node_encoder["U-Net encoder<br/>[MSHNet.py]"]
  node_middle["Middle layer<br/>[MSHNet.py]"]
  node_decoder["U-Net decoder<br/>[MSHNet.py]"]
  node_heads["Multi-scale heads<br/>[MSHNet.py]"]
end

subgraph group_dataeval["Data and evaluation"]
  node_dataset["IRSTD dataset<br/>[data.py]"]
  node_augment["Paired transforms<br/>[data.py]"]
  node_miou["mIoU<br/>[metric.py]"]
  node_pd_fa["PD/FA metric<br/>[metric.py]"]
  node_roc["ROC metric<br/>[metric.py]"]
end

subgraph group_objective["Training objective"]
  node_sls["SLS IoU loss<br/>[loss.py]"]
  node_location["Location penalty<br/>[loss.py]"]
end

node_user(("Researcher"))

node_user -->|"starts run"| node_runner
node_runner -->|"builds loaders"| node_dataset
node_dataset -->|"transforms pairs"| node_augment
node_runner -->|"constructs and calls"| node_mshnet
node_mshnet -->|"encodes"| node_encoder
node_mshnet -->|"processes bottleneck"| node_middle
node_mshnet -->|"decodes"| node_decoder
node_mshnet -->|"produces predictions"| node_heads
node_runner -->|"computes loss"| node_sls
node_sls -->|"adds location penalty"| node_location
node_runner -->|"updates weights"| node_optimizer
node_runner -->|"updates and reads"| node_miou
node_runner -->|"updates and reads"| node_pd_fa
node_runner -->|"updates and reads"| node_roc
node_runner -->|"loads and saves"| node_weights
node_runner -->|"prints or logs"| node_report
node_heads -->|"predictions scored"| node_sls
node_dataset -->|"supplies image batches"| node_mshnet
node_dataset -->|"supplies labels"| node_miou
node_dataset -->|"supplies labels"| node_pd_fa
node_dataset -->|"supplies labels"| node_roc
node_report -->|"reports results"| node_user

click node_runner "https://github.com/bit-ruiliu/mshnet/blob/main/main.py"
click node_dataset "https://github.com/bit-ruiliu/mshnet/blob/main/utils/data.py"
click node_augment "https://github.com/bit-ruiliu/mshnet/blob/main/utils/data.py"
click node_mshnet "https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py"
click node_encoder "https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py"
click node_middle "https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py"
click node_decoder "https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py"
click node_heads "https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py"
click node_sls "https://github.com/bit-ruiliu/mshnet/blob/main/model/loss.py"
click node_location "https://github.com/bit-ruiliu/mshnet/blob/main/model/loss.py"
click node_optimizer "https://github.com/bit-ruiliu/mshnet/blob/main/main.py"
click node_miou "https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py"
click node_pd_fa "https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py"
click node_roc "https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_runner,node_weights,node_report,node_optimizer toneBlue
class node_mshnet,node_encoder,node_middle,node_decoder,node_heads toneAmber
class node_dataset,node_augment,node_miou,node_pd_fa,node_roc toneMint
class node_sls,node_location toneRose
class node_user toneIndigo
```

## Components

### Run orchestration

- [Train/test runner](https://github.com/bit-ruiliu/mshnet/blob/main/main.py): `main.py`
- **Weights and checkpoints**
- **Metrics output**
- [Adagrad optimizer](https://github.com/bit-ruiliu/mshnet/blob/main/main.py): `main.py`

### Detector network

- [MSHNet](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py): `model/MSHNet.py`
- [U-Net encoder](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py): `model/MSHNet.py`
- [Middle layer](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py): `model/MSHNet.py`
- [U-Net decoder](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py): `model/MSHNet.py`
- [Multi-scale heads](https://github.com/bit-ruiliu/mshnet/blob/main/model/MSHNet.py): `model/MSHNet.py`

### Data and evaluation

- [IRSTD dataset](https://github.com/bit-ruiliu/mshnet/blob/main/utils/data.py): `utils/data.py`
- [Paired transforms](https://github.com/bit-ruiliu/mshnet/blob/main/utils/data.py): `utils/data.py`
- [mIoU](https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py): `utils/metric.py`
- [PD/FA metric](https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py): `utils/metric.py`
- [ROC metric](https://github.com/bit-ruiliu/mshnet/blob/main/utils/metric.py): `utils/metric.py`

### Training objective

- [SLS IoU loss](https://github.com/bit-ruiliu/mshnet/blob/main/model/loss.py): `model/loss.py`
- [Location penalty](https://github.com/bit-ruiliu/mshnet/blob/main/model/loss.py): `model/loss.py`

### Other components

- **Researcher**

## Connections

- Researcher → Train/test runner: starts run (evidence: `README.md`)
- Train/test runner → IRSTD dataset: builds loaders (evidence: `main.py`)
- IRSTD dataset → Paired transforms: transforms pairs (evidence: `utils/data.py`)
- Train/test runner → MSHNet: constructs and calls (evidence: `main.py`)
- MSHNet → U-Net encoder: encodes (evidence: `model/MSHNet.py`)
- MSHNet → Middle layer: processes bottleneck (evidence: `model/MSHNet.py`)
- MSHNet → U-Net decoder: decodes (evidence: `model/MSHNet.py`)
- MSHNet → Multi-scale heads: produces predictions (evidence: `model/MSHNet.py`)
- Train/test runner → SLS IoU loss: computes loss (evidence: `main.py`)
- SLS IoU loss → Location penalty: adds location penalty (evidence: `model/loss.py`)
- Train/test runner → Adagrad optimizer: updates weights (evidence: `main.py`)
- Train/test runner → mIoU: updates and reads (evidence: `main.py`)
- Train/test runner → PD/FA metric: updates and reads (evidence: `main.py`)
- Train/test runner → ROC metric: updates and reads (evidence: `main.py`)
- Train/test runner → Weights and checkpoints: loads and saves (evidence: `main.py`)
- Train/test runner → Metrics output: prints or logs (evidence: `main.py`)
- Multi-scale heads → SLS IoU loss: predictions scored (evidence: `main.py`)
- IRSTD dataset → MSHNet: supplies image batches (evidence: `main.py`)
- IRSTD dataset → mIoU: supplies labels (evidence: `main.py`)
- IRSTD dataset → PD/FA metric: supplies labels (evidence: `main.py`)
- IRSTD dataset → ROC metric: supplies labels (evidence: `main.py`)
- Metrics output → Researcher: reports results