A general framework for infrared small target detection and segmentation. By modifying or add the config file, you can adjust various parameters, switch models and datasets and so on, and you can easily add your own models and datasets and so on. It is recommended to spend a little time reading the tutorial before use, which can make you master the use of the framework faster.
Please refer to get_started.md for installation and dataset preparation.
python train.py <CONFIG_FILE>
For example, train ACM model with fpn in single gpu, run:
python train.py configs/acm/acm_res20_fpn_512x512_800e_nuaa.py
nproc_per_node is the number of gpus you are using.
python -m torch.distributed.launch --nproc_per_node=2 train.py <CONFIG_FILE>
For example, train ACM model with fpn and 2 gpus, run:
python -m torch.distributed.launch --nproc_per_node=2 train.py configs/acm/acm_res20_fpn_512x512_800e_nuaa.py
- You can specify the GPU at the second line of
os.environ['CUDA_VISIBLE_DEVICES']in train.py. - Be sure to set args.local_rank to 0 if using Multi-GPU training.
python test.py <CONFIG_FILE> <SEG_CHECKPOINT_FILE>
For example, test ACM model with fpn, run:
python test.py configs/acm/acm_res20_fpn_512x512_800e_nuaa.py work_dirs/acm_res20_fpn_512x512_800e_nuaa/20221009_231431/best.pth.tar
If you want to visualize the result, you only add --show at the end of the above command.
This part explains the config file in detail, which can make you understand the content of the config file more effectively and quickly, and master the overall framework.
Please refer to add_model.md for add custom model.
Please refer to add_dataset.md for add custom dataset.
Please refer to add_loss.md for add loss function.
Notice: Although you need to use the loss function already in pytorch, you still need to do this.
Please refer to add_optimizer.md for add optimizer and scheduler.
Notice: Although you need to use the optimizer and scheduler already in pytorch, you still need to do this.