Weakly Supervised Semantic Segmentation (WSSS) with only image-level supervision is a promising approach to deal with the need for Segmentation networks, especially for generating a large number of pixel-wise masks in a given dataset. However, most state-of-the-art image-level WSSS techniques lack an understanding of the geometric features embedded in the images since the network cannot derive any object boundary information from just image-level labels. We define a boundary here as the line separating the object and background. To address this drawback, we propose our novel ReFit framework, which deploys state-of-the-art class activation maps combined with various post-processing techniques in order to achieve fine-grained higher-accuracy segmentation masks. To achieve this, we investigate a state-of-the-art unsupervised semantic segmentation network that can be used to construct a boundary map, which enables ReFit to predict object locations with sharper boundaries. By applying our method to WSSS predictions, we were able to achieve up to 10% improvements even to the benefit of the current state-of-the-art WSSS methods for medical imaging.
- Dependencies :
matplotlib 3.5.2 numpy 1.21.5 Pillow 9.2.0 scikit-image 0.19.2 scikit-learn 1.0.2 scipy 1.9.1 torch 1.13.0 torchvision 0.14.0 nibabel 5.0.0
- The BraTS-2020 dataset can downloaded from this link.
- The preprocessed and 3-fold cross-validation split of prostate DECATHALON dataset WSS-CMER's link.
Basic dataset folder structure, using Prostate dataset as an exemplary. (Note: Make sure to change the dataset directory accordingly inside the config file )
Please set all paths as mentioned at the top of every program.
- Generate USS images
python deca_USS.py
- Train an image classifier for generating CAMs
python deca_Classifier.py
- Generate CAMs
python deca_GradCAM.py
- Refine CAMs with BoundaryFit module
python deca_BOUNDARY_FIT.py
- Evaluate the model
python deca_eval.py
Please set all paths as mentioned at the top of every program.
- Covert BraTS dataset
python brats_transformation.py
- Generate USS images
python brats_USS.py
- Train an image classifier for generating CAMs
python brats_Classifier.py
- Generate CAMs
python brats_GradCAM.py
- Refine CAMs with BoundaryFit module
python brats_BOUNDARY_FIT.py
- Evaluate the model
python brats_eval.py
- Create USS segementations
python3 BUSI_USS.py
python3 BUSI_USS.py --segment quick
- Generate CAMs
python3 Grad_cam.py
- Refine CAMs with BoundaryFit module
python3 BoundaryFit_busi.py
- Evaluate the model
python3 evaluate_busi.py
Qualitative segmentation results on BraTS and DECATHLON