526互联

Segment Anything（SAM）环境安装&代码调试

发布时间 2024-01-09 14:29:57作者: 要养家的程序猿

引子

　　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 ./
0
apt-get update && apt-get install libgl1
0
apt-get install libglib2.0-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
0
四、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)