引子
Segment Anything是前阵子大火的CV领域模型，之前也有尝试，只是没有整理。OK，让我们开始吧
一、拉取下载docker镜像
docker pull cnstark/pytorch:2.0.1-py3.9.17-cuda11.8.0-ubuntu20.04
二、安装SAM环境
docker run -it --gpus=all -v /datas/work/zzq:/workspace 8fd9e4c5e7bc bash
pip install opencv-python pycocotools matplotlib onnxruntime onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
git clone https://github.com/facebookresearch/segment-anything
cd /workspace/SAM/segment-anything
python scripts/amg.py --checkpoint sam_vit_h_4b8939.pth --input 170425986850.png --output ./
apt-get update && apt-get install libgl1
apt-get install libglib2.0-0
pip install segment_anything
python scripts/amg.py --checkpoint sam_vit_h_4b8939.pth --model-type vit_h --input 170425986850.png --output ./
三、可视化效果查看
python scripts/mask_generator.py
四、mask_generator.py代码
1 import numpy as np
2 import torch
3 import cv2
4
5 def apply_color_mask(image, mask, color, color_dark = 0.5):#对掩体进行赋予颜色
6 for c in range(3):
7 image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - color_dark) + color_dark * color[c], image[:, :, c])
8 return image
9
10
11
12
13 image_origin = cv2.imread('test_img/20240108094728.jpg')
14 image = cv2.cvtColor(image_origin, cv2.COLOR_BGR2RGB)
15 import sys
16 sys.path.append("..")
17 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
18
19 sam_checkpoint = "sam_vit_h_4b8939.pth"
20 model_type = "vit_h"
21
22 device = "cuda"
23
24 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
25 sam.to(device=device)
26
27 mask_generator = SamAutomaticMaskGenerator(sam)
28 masks = mask_generator.generate(image)
29
30 print(len(masks))
31 print(masks[0].keys())
32
33
34 image_select = image_origin.copy()
35 for i in range(len(masks)):
36 color = tuple(np.random.randint(0, 256, 3).tolist())#随机列表颜色，就是
37 selected_mask=masks[i]['segmentation']
38 selected_image = apply_color_mask(image_select,selected_mask, color)
39 cv2.imwrite("res.jpg", selected_image)