This is a pneumonia classification project that addresses the issue of class imbalance by utilizing generative adversarial networks (GAN) to generate images of minority class samples. In addition, the Spatial Attention Mechanism is introduced into ResNet18 to enhance the generalization performance of classifier!Moreover, this project adopts the FPGM pruning strategy to obtain a lightweight model!
🔥 Workflow
- ✅ Mar 21, 2023. Creat "MedGAN-ResLite" project repository and Find MedMNIST.
- ✅ Mar 22, 2023. Generate pneumonia samples with DCGAN.
- ✅ Mar 30, 2023. Replace original Loss function with Hinge Adversial Loss.
- ✅ Apri 1, 2023. DCGAN + Spectral Normalization.
- ✅ Apri 4, 2023. Add DCGAN metrics:Inception Score + FID + KID; Fuse and Split dataset;
- ✅ Apri 5, 2023. Override the dataset inheritance class.
- ✅ Apri 6, 2023. Write train, eval and infer scripts for classifier. And get a new-model by modifing input & output shape of pre-trained model. Add metrics:acc + auc + f1 + confusion matrix.
- ✅ April 7, 2023. Add scripts: export_onnx.py and inference_onnx.py.
- ✅ April 8, 2023. Tuning the hyperparameters of DCGAN.
- ✅ April 10, 2023. Explore CBAM attention mechanism to add location and quantity.
- ✅ April 14, 2023. Abalation Study: GAN, DCGAN, DCGAN+Hinge, DCGAN + SN, DCGAN + SH.
- ✅ April 21, 2023. Attention mechanism visualization using CAM.
- ✅ April 25, 2023. Make a Presentation.
- Coming Back!
- ✅ Mar 10, 2024. The dataset was preprocessed using Chest X-ray 2017 with reference to MedMNIST [paper] practices.
- ✅ Mar 11, 2024. Train GAN & CNN again!
- ✅ Mar 13, 2024. Histogram equalization was tried, but it did not work well~
- ✅ Mar 15, 2024. Attempts were made to introduce residual connection in GAN, but the effect was not good and the training speed was affected~
- ✅ Mar 20, 2024. Trying the WGAN training strategy and introducing Wasserstein distance did not work well~
- ✅ Mar 24, 2024. Add Pruning Sample by NNI.
- ✅ May 15, 2024. Release PulmoInsight Web Applicaiton!
- ✅ May 21, 2024. Release of MedGAN-ResLite-V2!
Clone repo and install requirements.txt.
git clone git@github.com:MaitreChen/MedGAN-ResLite.git
cd MedGAN-ResLite
pip install -r requirements.txt
You can preprocess the dataset by yourself, or you can get data_v2
directly from this link.
It includes the pneumoniamnist original real dataset and the fake dataset synthesized using GAN (see data README.md for preprocess and other details)
The dataset structure directory is as follows:
MedGAN-ResLite/
|__ data/
|__ real/
|__ train/
|__ normal/
|__ img_1.png
|__ ...
|__ pneumonia/
|__ img_1.png
|__ ...
|__ val/
|__ normal/
|__ pneumonia/
|__ test/
|__ ...
|__ fake/
|__ ...
You can download pretrained_v2
checkpoints from this link and put it in your pretrained/ folder. It contains resnet18-sam, sh-dcgan and other.(see README.md for details.)
🚀Quick start, and the results will be saved in the figures/classifier_torch folder.
python infer_classifier.py --ckpt-path pretrained/resnet18-sam.pth --image-path imgs/normal1.jpeg
🌜Here are the options in more detail:
Option | Description |
---|---|
--ckpt-path | Checkpoints path to load the pre-trained weights for inference. |
--image-path | Path of the input image for inference. |
--device | Alternative infer device, cpu or cuda, default is cpu. |
📛Note
If you want to visualize the attention mechanism, run the following command and the results will be saved in the figures/heatmap folder.
python utils/cam.py --image-path imgs/pneumonia_img1.png
💖More information about CAM can be found here!
🚀Quick start, and the results will be saved in the figures/generator_torch/single folder.
python infer_generator.py --ckpt-path pretrained/sh-dcgan.pth --batch-size 1 --mode -1
📛Note
If you want to generate fake images for training or sprite images, run following commands:
-
Generate a Sprite map. 【save results in figures/generator_torch/sprite】
python infer_generator.py --ckpt-path pretrained/sh-dcgan.pth --batch-size 64 --mode 1
-
Generate a batch of images. 【save results in figures/generator_torch/images】
python infer_generator.py --ckpt-path pretrained/sh-dcgan.pth --batch-size 64 --mode 0
💨When you generate a batch of images, batch-size is whatever you like❤
python eval_classifier.py --ckpt-path pretrained/resnet18-sam.pth
To evaluate a model, make sure you have torch-fidelity installed in requirements.txt❗
Then, you should prepare two datasets❗
- training datasets in data/merge folder. 【real images】(Note:merge "Normal" category data used for training from the
data/real
folder into thedata/merge
folder) - generation datasets in figures/generator_torch/image folder. 【fake images】
Everything is ready, you can execute the following command:
fidelity --gpu 0 --isc --input1 data/real_valid_normal_images --input2 figures/generator_torch/images
💖You can also set mode "--fid"
or "--kid"
. More information about fidelity can be found here!
python train_classifier.py
💝 More details about training your own dataset
Please refer to config/config.yaml and README.md.
python train_dcgan.py
If you want to export the ONNX model for ONNXRuntime or OpenVINO, please refer to README.md!
Method | Inception Score | FID | KID |
---|---|---|---|
GAN | 2.09 | 120.58 | 0.15 |
DCGAN | 2.09 | 92.50 | 0.11 |
SH-DCGAN | 2.09 | 36.92 | 0.03 |
Method | Inception Score | FID | KID |
---|---|---|---|
DCGAN | 2.09 | 92.50 | 0.11 |
DCGAN + Hinge | 2.09 | 68.54 | 0.06 |
DCGAN + SN | 2.09 | 45.68 | 0.04 |
SH-DCGAN | 2.09 | 36.92 | 0.03 |
Model | Accuracy | Precision | Recall | F1 score | Params | FLOPs | Inference Time |
---|---|---|---|---|---|---|---|
AlexNet | 90.16 | 90.16 | 90.16 | 90.16 | 14.57 | 270.01 | 0.02800 |
VGG16 | 91.22 | 92.23 | 91.22 | 91.17 | 27.56 | 15301.67 | 0.75499 |
VGG19 | 91.76 | 92.70 | 91.76 | 91.71 | 32.86 | 19463.47 | 1.06400 |
ResNet34 | 92.55 | 93.26 | 92.55 | 92.52 | 23.45 | 4008.46 | 0.32200 |
ResNet50 | 91.15 | 92.44 | 92.15 | 92.14 | 23.45 | 4008.46 | 0.30699 |
MobileNetV2 | 92.29 | 92.60 | 92.29 | 92.27 | 2.19 | 292.27 | 0.13999 |
ResNet18 | 92.15 | 93.30 | 90.13 | 91.34 | 11.16 | 1734.89 | 0.08696 |
ResNet18-SAM | 94.87 | 94.53 | 94.53 | 94.53 | 11.16 | 1734.89 | 0.09699 |
For any questions or suggestions about this project, welcome everyone to raise issues!
Also, please feel free to contact hbchenstu@outlook.com.
Thank you, wish you have a pleasant experience~~💓🧡💛💚💙💜