Exploring CNN and ViT for Weakly-Supervised Surgical Tool Segmentation
- [Pytorch]
- Some basic python packages such as Numpy, Scikit-image, SimpleITK, Scipy, Medpy ......
In this project, we use MICCAI Robotic Instrument Segmentation Challenge 2017 Official Link.
- Clone the repo:
git clone https://github.com/ziyangwang007/CV-WSL-Robot.git
cd CV-WSL-Robot
- Download the pre-processed data
Download the pre-processed data and put the data in ../data/robotic
. You can download the pre-processed dataset for Weakly-Supervised Learning study purpose, i.e. scribble annotation. Google Drive Google Drive Link, or Baidu Netdisk Baidu Netdisk Link with passcode: '8zf8'.
- (Optional) Generate Scribble Annotations
We generate scribble annotations via previous work code/scribbles_generator.py
. Please kindly check the code, the input is the pixel-level ground truth, and the output is the scribble annotation.
- Train the model
cd code
python train_Ours_Weakly_Consistency_Robot_2D.py
We provide some baseline methods as well.
Fully Supervised - CNN (UNet) -> Paper Link
python train_fully_supervised_2D
Fully Supervised - ViT (SwinUNet) -> Paper Link
python test_2D_fully_ViT
- Test the model
python test_2D_fully_ViT.py
or
python test_2D_fully.py
or
python test_2D_vit.py
TBC