CV-WSL-Robot

Exploring CNN and ViT for Weakly-Supervised Surgical Tool Segmentation

Demo Picture

Requirements

• [Pytorch]
• Some basic python packages such as Numpy, Scikit-image, SimpleITK, Scipy, Medpy ......

Datasets

In this project, we use MICCAI Robotic Instrument Segmentation Challenge 2017 Official Link.

Usage

1. Clone the repo:

git clone https://github.com/ziyangwang007/CV-WSL-Robot.git
cd CV-WSL-Robot

2. 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'.

3. (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.

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

5. Test the model

python test_2D_fully_ViT.py

or 

python test_2D_fully.py

or 

python test_2D_vit.py

Reference

TBC