Dense Nested Attention Network for Infrared Small Target Detection

In plain words

Infrared cameras often pick up a target that is only a few pixels wide, such as a ship or an aircraft far away. Common image networks shrink the picture several times on the way through, and a dot that small can disappear during that shrinking. Earlier learning methods mostly borrowed designs made for large objects, so they often lost these targets.

The paper's network keeps the shrinking, but it adds many connections between the stages so that details from every scale are passed back and forth. Attention steps then tell the network which channels and which locations matter. The output is a mask that marks where the targets are, and a small clustering step turns each blob of the mask into a single point, one per target.

The authors also built a training set of synthetic infrared images with exact labels, since real labelled data for this task is scarce.

TL;DR

• Task: single-frame infrared small target detection, output a mask and one centroid per target §3.1 §3.4.
• Method: a U-shaped encoder-decoder with densely nested skip connections (DNIM), a channel-then-spatial attention block (CSAM) after each fusion, and a pyramid fusion at the output §3.2 §3.3.
• Data: a new synthetic dataset, NUDT-SIRST, with 1327 images of 256x256 pixels §4 tab. I.
• Result (author-reported, NUDT-SIRST, ResNet-18 variant): best IoU and Pd in tab. II; the ResNet-34 variant has the lowest Fa.
• Status: the released code fills many gaps the paper leaves open, but several paper-code mismatches change the training recipe (see Paper vs code).

Problem

Targets in single-frame infrared images are small and dim. The paper says they range from one pixel to tens of pixels, have low signal-to-clutter ratio, and have little shape information §1. Pooling in a deep network can erase a one-pixel target before the output stage, so networks borrowed from generic object detection lose these targets in deep layers §1 §3.2.1. Traditional methods (filtering, local contrast, low-rank) depend on hand-set parameters and fail when target size or background changes §2.1. Pixel-level metrics such as IoU are unsuitable for targets of a few pixels, so the paper adds target-level metrics (Pd, Fa) §5.1. Open datasets are few and small, and existing ones have coarse or inaccurate labels §2.2.

Assumptions: input is a single infrared frame; the paper evaluates at 256x256 resolution §5.2.

Method

The network takes a 256x256 image and returns one logit per pixel. A clustering step then produces target centroids §3.1.

1. Preprocess and augment (training only). Normalise, then apply random flipping, blurring and cropping §5.2. The paper does not give the augmentation ranges; the released code does (see Hyperparameters).
2. Feature extraction (DNIM). Five down-sampling levels i = 0..4 form a grid of nodes x(i, j). The first column is an ordinary U-shaped encoder path: each level is max-pooled from the level above eq. 1. Each further column j receives, for node (i, j): all earlier nodes of the same level, the node one level deeper in column j-1 (upsampled), and the node one level shallower in column j (down-sampled) eq. 2 fig. 4. Nodes are concatenated before each convolution block.
3. Node block. Each node applies two 3x3 convolution stages with batch norm, then CSAM §3.2.3:
   • channel attention eq. 3-eq. 4: sigmoid of a shared one-hidden-layer MLP applied to the global max and global average of each channel, then multiply;
   • spatial attention eq. 5-eq. 6: a 7x7 convolution on the channel-wise max and mean of the feature map, sigmoid, then multiply. The released code adds a residual shortcut and a final ReLU; the paper's formulas do not show them model_DNANet.py.
4. Output fusion (FPFM). The top-level nodes of the five levels are projected to a common channel count, upsampled to full resolution, concatenated, and refined by one block eq. 7 §3.3 fig. 3. The paper writes this as a concatenation of the top-row nodes of each level.
5. Target extraction. Threshold the output (see Results for the operating point), label 8-connected components of foreground pixels, and take each component's centroid eq. 8 eq. 9 §3.4. Two pixels belong to one component when their 8-neighbourhoods intersect and they share the same binary value.
6. Loss. Soft-IoU on the sigmoid of the output, optimised with Adagrad §5.2.
7. Evaluation. IoU at pixel level, Pd and Fa at target level with a 3-pixel centroid tolerance eq. 11 eq. 12 eq. 13 §5.1.

# Nested node grid  [eq. 1-2]  [code: model/model_DNANet.py#L103-L163]
# image: (B, 3, 256, 256); node[i][j] is the feature map at level i, column j

def build_nodes(image):
    # Column 0: plain encoder path, max-pool between levels  [eq. 1]
    node[0][0] = node_block(image)
    for level in range(1, 5):
        node[level][0] = node_block(max_pool(node[level - 1][0]))
    # Columns j >= 1: nodes with i + j <= 4  [eq. 2]
    for column in range(1, 5):
        for level in range(0, 5 - column):
            # Same-level nodes to the left, in the same column order
            same_level = [node[level][k] for k in range(column)]
            # Node one level deeper, one column left: upsampled x2
            deeper = upsample(node[level + 1][column - 1], scale=2)
            inputs = same_level + [deeper]
            # Node one level shallower, same column: halved in size
            if level > 0:
                shallower = downsample(node[level - 1][column], scale=0.5)
                inputs = inputs + [shallower]
            node[level][column] = node_block(concat(inputs))
    return node

Note: the code resamples with bilinear interpolation, and the paper writes max-pooling for the deeper input and upsampling for the shallower one. See Paper vs code.

# Node block: two conv stages, then CSAM, then residual  [§3.2.3] [eq. 3-6]
# [code: model/model_DNANet.py#L57-L88]

def node_block(inputs):
    # Two 3x3 conv + batch-norm stages  [§3.2.2]
    hidden = relu(batch_norm(conv3x3(inputs)))
    hidden = batch_norm(conv3x3(hidden))
    # Channel attention, then spatial attention  [eq. 3-4]
    channel_mask = sigmoid(mlp(global_max(hidden)) + mlp(global_avg(hidden)))
    hidden = channel_mask * hidden
    # Spatial attention with a 7x7 conv on channel-wise max and mean  [eq. 5-6]
    spatial_mask = sigmoid(conv7x7(concat(max_channels(hidden),
                                          mean_channels(hidden))))
    hidden = spatial_mask * hidden
    # Residual shortcut and final ReLU: in code, not in the paper's formulas
    shortcut = batch_norm(conv1x1(inputs))   # projection when channels differ
    return relu(hidden + shortcut)

# Output fusion and deep supervision  [eq. 7] [§3.3]
# [code: model/model_DNANet.py#L125-L136, #L165-L177]

def fuse_outputs(node):
    # Top-row nodes: node[i][4 - i] for i = 0..4 (J is undefined in the paper)
    scaled = []
    for level in range(0, 5):
        projected = conv1x1(node[level][4 - level])   # 16 channels
        scaled.append(upsample(projected, scale=2 ** level))
    fused = node_block(concat(scaled + [node[0][4]]))
    # Four heads in training: three side outputs and the fused output
    heads = [conv1x1(node[0][1]), conv1x1(node[0][2]),
             conv1x1(node[0][3]), conv1x1(fused)]
    return heads

# Soft-IoU loss on one batch  [§5.2]  [code: model/loss.py#L5-L23]
# logits: (B, 1, H, W); target: (B, 1, H, W), values in {0, 1}

def soft_iou_loss(logits, target):
    # Foreground probability per pixel  [code: model/loss.py#L7]
    probs = sigmoid(logits)
    # Soft intersection and union summed over the whole batch  [code: #L13-L14]
    intersection = sum(probs * target)
    union = sum(probs) + sum(target) - intersection
    # Smoothing constant 1 in numerator and denominator  [code: #L8, #L20]
    return 1 - (intersection + 1) / (union + 1)

# Training step with deep supervision  [§5.2]  [code: train.py#L72-L86]
# Adagrad, learning rate 0.05, batch 16 per the paper

def training_step(images, masks):
    # Average the Soft-IoU over all four heads  [code: train.py#L75-L80]
    heads = model(images)
    loss = mean([soft_iou_loss(head, masks) for head in heads])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss

# Target-level Pd and pixel-level Fa for one image  [eq. 12-13]
# [code: model/metric.py#L52-L108]

def detection_counts(pred_mask, gt_mask):
    # 8-connected components  [§3.4] [code: metric.py#L71-L74]
    pred_parts = components(pred_mask, connectivity=8)
    gt_parts = components(gt_mask, connectivity=8)
    matched_areas = []
    # Greedy match: each target takes the first unused predicted centroid
    # within 3 pixels, in list order, not the nearest
    # [code: metric.py#L86-L97]
    for target in gt_parts:
        for part in pred_parts:
            if distance(part.centroid, target.centroid) < 3:
                matched_areas.append(part.area)
                pred_parts.remove(part)
                break
    # Matched count gives Pd; unmatched predicted pixels give Fa
    # Unmatched is judged by area value, not identity  [code: metric.py#L99]
    false_pixels = sum(part.area for part in pred_parts)
    return len(matched_areas), len(gt_parts), false_pixels

Pd over the dataset is matched targets divided by all targets eq. 12. Fa over the dataset is false pixels divided by 256 x 256 x number of images eq. 13 [code: metric.py#L105-L106]. The code hard-codes 256x256 metric.py.

Synthetic NUDT-SIRST recipe §4

The dataset is generated, not recorded. Recipe, as stated:

1. Point, spot and extended targets are simulated from a Gaussian kernel and from collected target templates (spot, plane, ship, UAV) §4.2 fig. 7.
2. A scene-aware CNN classifies each background among the five scene types (city, field, highlight, sea, cloud) §4.1 §4.2. A pre-defined probability selects target size and type for that background §4.2.
3. Target intensity is set so the signal-to-clutter ratio equals a value C drawn from {3, 4, 5, 6} eq. 10 §4.2.
4. A 5x5 Gaussian blur with sigma in {0.2, 0.5, 1.0, ...} smooths the boundary §4.2.
5. Low-quality images are removed by hand §4.2.

Most backgrounds were collected by the authors; a few field backgrounds come from a cited external source §4.1. The repository has no generation code and the paper does not give the probability values or the full sigma list.

Hyperparameters

Results

Author-reported. Setup: NUAA-SIRST and NUDT-SIRST, train-to-test ratio 1 §5.2; 256x256 inputs; all CNN baselines retrained on the same split with the authors' own re-implementations §5.3. Units: IoU x10^-2, Pd x10^-2, Fa x10^-6 tab. II. Highlights in the paper: red is best, blue is second best.

Headline numbers (DNA-Net-ResNet-18, from tab. II):

• NUDT-SIRST: IoU 87.09, Pd 98.73, Fa 4.223. Best IoU and Pd. Fa is second to DNA-Net-ResNet-34 (3.710).
• NUAA-SIRST: IoU 77.47, Pd 98.48, Fa 2.353. Pd and Fa are best; IoU is second to DNA-Net-ResNet-34 (77.54).

Table II (main comparison, Tr = 50%):

Source: tab. II. Traditional methods are listed first in the paper, then CNN methods, then this paper's variants.

Computational cost on NUDT-SIRST tab. III:

Pd / Fa under different centroid tolerances on NUDT-SIRST tab. IV:

Ablation: dense nested skip module (DNIM) tab. V:

Ablation: CSAM tab. VI:

Ablation: feature pyramid fusion output (FPFM) tab. VII:

Training on a mix of real (NUAA) and synthetic (NUDT) images tab. VIII; the paper's own counts are kept as printed:

Dataset comparison tab. I:

Qualitative results and the ROC curves are in fig. 8, fig. 9 and fig. 12; they are described in words only in this card.

Paper vs code

Official repository: the paper's footnote and conclusion link github.com/YeRen123455/Infrared-Small-Target-Detection as the authors' repository §5.3. Checked at commit 21584be.

• Batch size. Paper: 16 §5.2. Code default: 4 parse_args_train.py. The README's training command passes 16 README.md.
• Learning-rate schedule. Paper names CosineAnnealingLR §5.2. The code calls scheduler.step() once, before training train.py; no other call exists in the repo, so the learning rate stays at 0.05 for all 1500 epochs as released.
• Down/up-sampling in the nest. Paper eq. 2 applies max-pooling to the deeper input and upsampling to the shallower one eq. 2. Code upsamples the deeper input by 2 and halves the shallower input with bilinear interpolation model_DNANet.py model_DNANet.py.
• Pooling in CSAM. Paper text says the pooling in the attention blocks uses stride 2 §3.2.3. Code uses global pooling for channel attention model_DNANet.py and channel-wise max and mean for spatial attention model_DNANet.py.
• Residual path. Code adds a residual shortcut and a final ReLU to each node block model_DNANet.py. The paper's formulas do not include them.
• Channel-attention width. Code uses a hidden width of C/16 model_DNANet.py. The paper says "one hidden layer" and gives no ratio.
• Output notation J. Paper eq. 7 uses an index J without a value. Code uses the top-row nodes x(i, 4-i) model_DNANet.py.
• Deep supervision. Not in the paper. Code trains with four heads and averages the Soft-IoU over them train.py model_DNANet.py. Test uses the fused head only test.py.
• Residual blocks per level. Not in the paper. Level 0 always has one block; level i (i >= 1) has num_blocks[i-1] blocks model_DNANet.py.
• VGG-10. Code's vgg_10 config uses the same residual CBAM block as the ResNets load_param_data.py train.py. The class VGG_CBAM_Block is defined but not used model_DNANet.py.
• Xavier init. Code applies Xavier normal to convolution weights only; biases and batch-norm parameters keep their defaults utils.py train.py. Paper says "weights and bias" §5.2.
• Normalisation. Code uses ImageNet mean and standard deviation train.py. Paper says only "normalized".
• Soft-IoU scope. Code sums intersection and union over the whole batch, not per image loss.py. Smoothing is 1 loss.py; the paper gives no formula.
• Threshold for Pd and Fa. The paper does not say which operating point its Pd and Fa come from. Code's PD_FA sweeps 11 thresholds from 0 to 255 applied to raw logits metric.py and saves the whole curve test.py. The IoU path uses logit > 0 metric.py.
• Dataset-level IoU. Code pools intersection and union over the whole test set metric.py metric.py. Paper eq. 11 does not say whether IoU is per image or pooled.
• Test split. Code evaluates on the test split after every epoch train.py. self.best_iou is set once and never updated train.py, so save_model writes the checkpoint each epoch to one filename utils.py. The file holds the last epoch, not necessarily the best one. The paper does not say how its model was chosen.
• Synthesis data. Paper: mixing experiment uses 791 synthetic training images tab. VIII. The released NUDT 50/50 train split has 663 ids train.txt. The paper does not reconcile these. No synthesis code is in the repo.
• Internal numbers in the paper. Table IV gives DNA-Net-ResNet18 Fa 4.228 and ALCNet Fa 9.581 at < 4 px; Table II gives 4.223 and 9.261. Table VIII's real-only row gives Fa 4.223 for NUAA, but Table II gives 2.353 for NUAA and 4.223 for NUDT.
• README results. The README lists a second DNANet-ResNet-18 row with released weights: NUDT 88.61 / 98.42 / 4.30 and NUAA 79.26 / 98.48 / 2.30 README.md. These are not in the paper.
• License file. The repo's LICENSE is MIT, but its copyright line names a different person (Jiwoon Ahn, 2018) LICENSE. Recorded here as MIT per the file text.

Pitfalls

• Section and table numbers. This card cites sections by the LaTeX counter (for example §3.2 is "III-B" in IEEE's rendering). Tables are cited with Roman numerals as IEEE prints them; figures and equations use the paper's own Arabic numbers.
• Eq. 2 cannot be implemented as written if the level index follows eq. 1: it upsamples the shallower node and max-pools the deeper one. Use the code's direction (upsample deeper, halve shallower) and note that the code uses bilinear resampling, not max-pool.
• Eq. 7 index J is undefined. Taking the top-row nodes x(i, 4-i) is the code's choice.
• Batch and schedule. To reproduce the paper's numbers you need batch 16 and a learning rate that actually decays. The released loop never decays it. With batch 16 and drop_last=True, 663 training images give 41 steps per epoch and 7 images are skipped each epoch train.py.
• Adagrad at lr 0.05. The paper gives only the lr. Any other Adagrad option is a library default.
• Soft-IoU is batch-level. A per-image version would give different gradients, especially for images with no target.
• Pd and Fa matching is fragile. Matching is greedy in list order, not nearest-first metric.py. Unmatched predicted components are identified by area value metric.py. Two components of equal area can be confused. The code hard-codes 256 x 256 metric.py.
• Fa normalisation. Fa divides by the full image area times the number of images metric.py. Units in the tables are x10^-6.
• Units in the tables. IoU and Pd are printed x10^-2 and Fa x10^-6 tab. II. Do not compare them to the x100 or raw values without converting.
• Model selection. The released code's saved checkpoint is the last epoch, not the best-IoU epoch (see Paper vs code). If you compare against the reported numbers, record which epoch you evaluated.
• Baselines are re-implementations. The authors re-trained every CNN baseline on their own split using their own code §5.3. The numbers are not the baselines' published numbers. Traditional baselines use the adaptive threshold of eq. 14; CNN baselines use fixed thresholds §5.3.
• Synthetic data recipe is incomplete. The probabilities that pick target size and type, the full set of sigma values and the criteria for "low quality" are not given. The scene-aware CNN is not released.
• Mixing table. The 791 count and the inconsistent denominators in tab. VIII make it hard to reproduce that experiment exactly.
• Cost. Table III reports train and inference times only for the light ResNet-10 variant and three baselines, which use ResNet-20 backbones tab. III. Times for DNA-Net-ResNet-18 and ResNet-34 are not reported. Timings were measured on the GeForce 3090 §5.2.
• NUST-SIRST is excluded from the main table. The paper says only about 30% of its targets meet the small-target definition, and the updated results are only in the authors' repository §5.3.

Lineage

• Builds on: U-shaped encoder-decoder networks with skip connections, UNet++-style nested skips, and CBAM-style channel and spatial attention, all as named in the paper's introduction §1.
• Compared against: traditional methods (Top-Hat, Max-Median, WSLCM, TLLCM, IPI, NRAM, RIPT, PSTNN, MSLSTIPT) and CNN methods (MDvsFA-cGAN, ACM, ALCNet) tab. II.
• Superseded by: not stated.