jittor2023-赛题二-SCU_MILab

简介

本项目包含了第三届计图挑战赛 - 大规模无监督语义分割赛题的代码实现。主要步骤为：

语义表征：1）参考PASS通过自监督的表征学习方法训练随机初始化的模型，以学习形状和类别表征。这里通过增加MLP隐层维度提高模型学习能力；2）通过基于像素注意力的聚类方法获得伪类别，并将生成的类别分配给每个图像像素，获得初始伪标签。
形状增强：1）使用伪标签的roi信息生成对应的point prompts，使用SAM（vit_b）模型对图像进行分割推理得到类别不可知的分割masks。2）由于SAM分割粒度未知，需要对SAM分割得到的masks与伪标签进行匹配；参考Personalize SAM with one shot，我们通过对每个mask的下image encoder的embedding与伪标签下image encoder的embedding进行相似度匹配并设置threshold，高于该阈值的即认为是目标类别mask，最终融合多个符合条件的mask得到fused mask；3）进一步使用fused mask生成point prompts，box prompts以及mask prompts通过SAM模型进行两次边缘微调，得到最后的伪标签。
不完全标注的伪标签学习：对于以上得到的伪标签，我们假设已标注的像素点都是positive label，未标注的则是negativa label or positive label。参考ShapePU，我们采用EM算法估计每个类别的标注比率，然后通过标注比率区分正负样本，接着最大化估计得到的负样本的边缘概率进行负样本监督，另一边对已标注的正样本进行正样本监督，以此对模型进行微调，提高分割质量。
推理：使用步骤3微调后的模型进行推理，将生成的标签分配给图像的每个像素。对于得到的分割结果，进一步使用步骤2进行形状微调，得到最终分割结果。
安装
PASS模型的环境要求：

System: Linux(e.g. Ubuntu/CentOS/Arch), macOS, or Windows Subsystem of Linux (WSL)
Python version >= 3.7
CPU compiler (require at least one of the following)
- g++ (>=5.4.0)
- clang (>=8.0)

第一步: 安装计图

计图的安装可以参考以下文档Jittor install

第二步: 安装依赖

python -m pip install scikit-learn
python -m pip install pandas
python -m pip install numpy
python -m pip install munkres
python -m pip install tqdm
python -m pip install pillow
python -m pip install opencv-python
python -m pip install faiss-gpu

训练

sh train.sh

测试

sh test.sh

CUDA_VISIBLE_DEVICES='0' python test.py \
--input_path ./test_img \
--output_path ./result

SAM权重转换

我们发现直接使用jittor 加载SAM模型进行推理，有较大的精度损失，这里我们将pytorch版本的SAM（vit-b）权重转换为jittor版本，并存储为pickle格式。

python torch2jt.py

模型checkpoint及SAM权重下载

可通过百度云盘下载。

引用

@article{gao2022luss,
  title={Large-scale Unsupervised Semantic Segmentation},
  author={Gao, Shanghua and Li, Zhong-Yu and Yang, Ming-Hsuan and Cheng, Ming-Ming and Han, Junwei and Torr, Philip},
  journal=TPAMI,
  year={2022}
}

@article{kirillov2023segany,
  title={Segment Anything},
  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{```
a}r, Piotr and Girshick, Ross},
  journal={arXiv:2304.02643},
  year={2023}
}

@article{zhang2023personalize,
  title={Personalize Segment Anything Model with One Shot},
  author={Zhang, Renrui and Jiang, Zhengkai and Guo, Ziyu and Yan, Shilin and Pan, Junting and Dong, Hao and Gao, Peng and Li, Hongsheng},
  journal={arXiv preprint arXiv:2305.03048},
  year={2023}
}

@inproceedings{zhang2022shapepu,
  title={Shapepu: A new pu learning framework regularized by global consistency for scribble supervised cardiac segmentation},
  author={Zhang, Ke and Zhuang, Xiahai},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={162--172},
  year={2022},
  organization={Springer}
}