Awesome

P²Net

Share us a :star: if this repo does help

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.

Updates

• 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.

Table of Contents

• Introduction
• Results and Trained Model
• Preliminaries
• Environments
• Run Details
  • Train Process
  • Test Process
  • Visualization
• Citation
• License

Introduction

We focus on the "imbalanced fineness" and "imbalanced appearances" in the fine-grained ship classification task in remote sensing images and propose P²Net to address the above issues. P²Net is a weakly-supervised method, and can be trained in an end-to-end way. Our P²Net can outperform many recent methods in the FGSC task.

Results and Trained Model

Models trained on augmented train dataset

Models trained on original train dataset without pre-augmentations

Preliminaries

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 P²Net 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]

Environments

• 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.)

Run Details

Train Process

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.

Test Process

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

Visualization

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.

Citation

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}}

License

This project is licensed under the MIT License. See LICENSE for details