HQ-SAM代码复现

代码地址SysCV/sam-hq: Segment Anything in High Quality [NeurIPS 2023] (github.com)

一：快速安装

适用于已经配好环境，可以直接上手使用。对于新手不推荐使用这种方式，不太容易看懂。

pip install segment-anything-hq
python
from segment_anything_hq import sam_model_registry
model_type = "<model_type>" #"vit_l/vit_b/vit_h/vit_tiny"
sam_checkpoint = "<path/to/checkpoint>"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)

from segment_anything_hq import sam_model_registry
model_type = "<model_type>" #"vit_l/vit_b/vit_h/vit_tiny"
sam_checkpoint = "<path/to/checkpoint>"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)

这段代码的意思是找到这个 segment_anything_hq 所在的位置，然后修改下面模型的名字和模型的地址。 

可以通过运行下面这个命令具体使用用例

export PYTHONPATH=$(pwd)
python demo/demo_hqsam_pip_example.py

二：标准安装

2.1环境配置

需要的环境有PyTorch和TorchVision依赖项，安装支持CUDA的PyTorch和TorchVision。

PyTorch官网地址Start Locally | PyTorch

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

具体安装Pytorch可以查看这篇文章在PyCharm中安装Pytorch的GPU版本-优快云博客

我这边做一个简单的介绍，在官网中找到符合自己型号的cuda，然后复制下面的命令到PyCharm中去下载。

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

2.2在本地克隆储存库并使用

git clone https://github.com/SysCV/sam-hq.git

cd sam-hq; pip install -e .

2.3安装依赖项

用于掩码后处理，以COCO格式保存掩码、示例笔记本以及ONNX格式导出模型是必需的。juptyer也需要运行notebooks。

pip install opencv-python pycocotools matplotlib onnxruntime onnx timm

2.4conda环境的创建

conda create --name sam_hq python=3.8 -y
conda activate sam_hq
conda install pytorch==1.10.0 torchvision==0.11.0 cudatoolkit=11.1 -c pytorch -c nvidia
pip install opencv-python pycocotools matplotlib onnxruntime onnx timm


# under your working directory
git clone https://github.com/SysCV/sam-hq.git
cd sam-hq
pip install -e .
export PYTHONPATH=$(pwd)

2.5模型检查点

HQ-SAM模型有三种不同的型号版本，对应不用的主干尺寸。这些模型可以直接使用。

我们首先先创建一个文件夹，文件夹就按这下面的命令创建。然后把下载的模型放到这个文件夹中。

mkdir pretrained_checkpoint

然后下载模型 ，由下面这张图我们可以看出，由B到L，再到H，模型的大小是越来越来大的，对应的推理速度也是越来越慢。

from segment_anything import sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")

上面这段代码的意思是，找到segment_anything所在的模块，然后修改下面模型的名称和模型的位置。模型在以下这个位置。

修改下面这两个地方。

2.6开始使用

配置完环境，修改完路径和名称就可以使用了。

要想获得HQ-SAM的视觉效果

python demo/demo_hqsam.py