KFGMoE

Keyframe-Guided Mixture of Experts for Multi-Domain Moving Infrared Small Target Detection

PyTorch implementation of KFGMoE. The manuscript is currently under review.

Introduction

KFGMoE is a multi-domain moving infrared small target detection framework in which the current keyframe guides expert selection and historical feature construction. It combines prototype-guided keyframe enhancement, MoE-based feature construction and spatio-temporal feature refinement in one jointly trained model. Spatial supervision connects detection-relevant keyframe regions to expert allocation, and temporal fusion combines the resulting features for current-frame detection.

Prerequisite

• Python 3.10.20
• PyTorch 2.9.0+cu128
• torchvision 0.24.0+cu128
• NumPy 2.1.3
• Pillow 11.1.0
• pycocotools 2.0.10
• tqdm 4.67.3
• Tested on Ubuntu 22.04.5 with CUDA 12.8 and one NVIDIA GeForce RTX 5090 (32 GB).

Install PyTorch and torchvision for your CUDA environment, then install the remaining dependencies:

python -m pip install -r requirements.txt
python tools/prepare_upstream.py

The second command prepares the required third-party modules; see THIRD_PARTY.md for details.

Datasets

Download the bounding-box versions of DAUB-H, IRDST-R and ITSDT-15K from the dataset download page.

Annotation lists use JSON Lines, with one keyframe per line:

{"path":"DAUB-H/train/data1/12.bmp","domain":1,"boxes":[[100,80,106,86,0]]}
{"path":"DAUB-H/train/data1/13.bmp","domain":1,"boxes":[]}

path is relative to --data-root; boxes use [x1, y1, x2, y2, 0] in original-image pixels, or [] for an empty frame. Training requires integer box coordinates.

Dataset	Domain ID
DAUB-H	1
IRDST-R	2
ITSDT-15K	3

Each video must contain consecutive BMP frames starting at 0.bmp, with indices matching the annotations. For one-based originals such as 001.bmp, prepare a separate zero-based copy aligned with the benchmark labels.

Convert an existing plain-text annotation file (image_path x1,y1,x2,y2,0 ...) with:

python tools/prepare_annotations.py --input original_train.txt --output annotations/train.jsonl --data-root /path/to/datasets

Repeat for test.jsonl. Use --strip-prefix /old/dataset/root for old absolute paths and --domain to specify a domain ID when it cannot be inferred from the path.

Usage

Train

Save the pretrained initialization weights as weights/pretrained.pth:

python train.py --config configs/kfgmoe.json --train-list annotations/train.jsonl --data-root /path/to/datasets --initial-weights weights/pretrained.pth --output outputs/kfgmoe

Settings are in configs/kfgmoe.json. To resume, replace --initial-weights weights/pretrained.pth with --resume outputs/kfgmoe/latest.pt.

Test

Save the trained KFGMoE weights as weights/kfgmoe_epoch060_ema.pth:

python test.py --config configs/kfgmoe.json --weights weights/kfgmoe_epoch060_ema.pth --annotations annotations/test.jsonl --data-root /path/to/datasets --output runs/test --device cuda

Visualization

Predict a video sequence:

python predict.py --weights weights/kfgmoe_epoch060_ema.pth --sequence /path/to/video/frames --domain DAUB-H --output runs/demo --device cuda

DAUB-H example from Fig. 6(a) of our manuscript. Green boxes indicate ground truth or correct detections, dashed blue boxes mark missed targets and yellow markers identify regions enlarged by ×3.

Results

Joint training on three domains:

Dataset	mAP50 (%)	Pr (%)	Rc (%)	F1 (%)	mAP75 (%)	mAP50:95 (%)	Fa ↓
DAUB-H	53.7530	70.4940	77.6877	73.9162	38.6289	33.8694	0.151322
IRDST-R	74.2099	83.2310	90.6922	86.8015	0.0126	14.4536	0.064314
ITSDT-15K	84.0296	90.2328	94.1169	92.1339	11.6903	31.3551	0.301000
Macro average	70.6642	81.3193	87.4989	84.2839	16.7773	26.5594	0.172212

At IoU 0.5, Rc is maximum recall, Pr is mean interpolated precision over that recall range and F1 is their harmonic mean. Fa is false detections per frame at confidence 0.5 and matching IoU 0.5. Macro averages weight the three domains equally.

Acknowledgements

We thank the CoMoE authors and dataset providers for sharing code, pretrained weights and data. See THIRD_PARTY.md for implementation sources.

Contact

Jinwu Lyu: jinwulyu@outlook.com.