The official implementation of "Contrastive Learning for Fine-grained Ship Classification in Remote Sensing Images". The paper can be accessed in [IEEE | Lab Server | ResearchGate]. (Accepted by TGRS 2022)
If you encounter any question, please feel free to contact us. You can create an issue or just send email to me windvchen@gmail.com. Also welcome for any idea exchange and discussion.
-
2023/07/22: We have added Google Drive links for all the datasets we used.
-
2023/03/30: To facilitate reproduction of our results and ensure the fairness of possible future comparative experiments, we have provided all of our split and augmented datasets in Preliminaries.
-
2022/07/16: The repository is now public.
-
2022/06/09: The code cleanup is finished and the complete codes are provided, also the weights of our model for FGSC-23 and FGSCR-42 datasets.
We focus on the "imbalanced fineness" and "imbalanced appearances" in the fine-grained ship classification task in remote sensing images and propose P2Net to address the above issues. P2Net is a weakly-supervised method, and can be trained in an end-to-end way. Our P2Net can outperform many recent methods in the FGSC task.
Method | Params(M) | FLOPs(G) | FGSC-23 (AA) | FGSCR-42(AA) |
---|---|---|---|---|
ResNet-50 | 23.6 | 4.12 | 86.92 | 91.62 |
HBPNet``ECCV18 | 74.9 | 6.59 | 87.72 | 91.32 |
DCL``CVPR19 | 23.8 | 4.12 | 85.35 | 90.24 |
TASN``CVPR19 | 34.8 | 18.7 | 87.03 | 91.85 |
GFNet``NIPS20 | 56.5 | 4.59 | 87.13 | 92.03 |
API-Net``AAAI20 | 23.6 | 4.12 | 87.78 | 91.47 |
ProtoTree``CVPR21 | 108.8 | 20.7 | 84.17 | 89.92 |
P2Net | 26.9 | 4.23 | 88.99 [Google Drive | Baidu Pan (code:tr3i)] |
93.21 [Google Drive | Baidu Pan (code:nyro)] |
Method | FGSC-23 (AA) | FGSCR-42(AA) |
---|---|---|
ResNet-50 | 85.68 | 91.85 |
HBPNet``ECCV18 | 86.09 | 92.09 |
DCL``CVPR19 | 84.31 | 90.65 |
TASN``CVPR19 | 86.11 | 92.87 |
GFNet``NIPS20 | 85.37 | 92.85 |
API-Net``AAAI20 | 85.32 | 91.92 |
ProtoTree``CVPR21 | 80.46 | 79.14 |
P2Net | 88.56 [Google Drive | Baidu Pan (code:4qxh)] |
94.19 [Google Drive | Baidu Pan (code:v06g)] |
Please at first download datasets FGSC-23 [code:n8ra] or FGSCR-42, then prepare the datasets as the following structure:
├── train
├── cls 1 # The first class
├── img_1.jpg # The first image, for FGSCR-42 it should be '.bmp'
├── img_2.jpg
├── ...
├── cls 2
├── ...
├── valid
├── test
We provide DatasetSplit.py to easily arrange the dataset into the above structure.
To overcome the imbalanced sample issue of the two datasets, you can make use of customTransform.py to augment the fewer-sample subclass in the train dataset. (Please notice that the val/test dataset is not suggested to be augmented, and you can also choose not to augment the train dataset to see the performance of P2Net on the imbalanced sample issue.)
You can also find here all the datasets we used for reproduction or comparative experiments: [Baidu Pan (code:4j35) | Google Drive]
- Windows/Linux both support
- python 3.8
- PyTorch 1.9.0
- torchvision
- pretrainedmodels
- wandb (Suggested, a good tool to visualize the training process. If not want to use it, you should comment out the related codes.)
To train our P^2Net
, run:
# For FGSC-23 dataset
python PPNet.py --data FGSC23_path --epochs 100 --workers 4 --batch-size 64 --num_classes 23 --proxy_per_cls 3 --pretrained
# For FGSCR-42 dataset
python PPNet.py --data FGSCR242_path --epochs 100 --workers 4 --batch-size 64 --num_classes 42 --proxy_per_cls 2 --pretrained
To train the baseline, just add --baseline
, run:
# For FGSC-23 dataset
python PPNet.py --data FGSC23_path --epochs 100 --workers 4 --batch-size 64 --num_classes 23 --pretrained --baseline
# For FGSCR-42 dataset
python PPNet.py --data FGSCR242_path --epochs 100 --workers 4 --batch-size 64 --num_classes 42 --pretrained --baseline
The architecture is ResNet-50 by default, and you can also change --arch
to other architecture like "densenet121"
or so on by your preference. The train results will be saved in model/expX.
Just add --test
and determine --test_model_path
when to conduct an inference.
To test our P^2Net
, run:
# For FGSC-23 dataset
python PPNet.py --data FGSC23_path --epochs 100 --workers 4 --batch-size 64 --num_classes 23 --proxy_per_cls 3 --test --test_model_path model_path
# For FGSCR-42 dataset
python PPNet.py --data FGSCR42_path --epochs 100 --workers 4 --batch-size 64 --num_classes 42 --proxy_per_cls 2 --test --test_model_path model_path
To test the baseline, run:
# For FGSC-23 dataset
python PPNet.py --data FGSC23_path --epochs 100 --workers 4 --batch-size 64 --num_classes 23 --baseline --test --test_model_path model_path
# For FGSCR-42 dataset
python PPNet.py --data FGSCR42_path --epochs 100 --workers 4 --batch-size 64 --num_classes 42 --baseline --test --test_model_path model_path
We provide two visualization tools drawCM.py and tsne.py. You can make use of them and modify PPNet.py a bit to visualize the confusion matrix or t-SNE results.
If you find this paper useful in your research, please consider citing:
@ARTICLE{9832938,
author={Chen, Jianqi and Chen, Keyan and Chen, Hao and Li, Wenyuan and Zou, Zhengxia and Shi, Zhenwei},
journal={IEEE Transactions on Geoscience and Remote Sensing},
title={Contrastive Learning for Fine-Grained Ship Classification in Remote Sensing Images},
year={2022},
volume={60},
number={},
pages={1-16},
doi={10.1109/TGRS.2022.3192256}}
This project is licensed under the MIT License. See LICENSE for details