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

引子

　　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 Simple Index

git clone GitHub - facebookresearch/segment-anything: The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.

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代码

import numpy as np
import torch
import cv2

def apply_color_mask(image, mask, color, color_dark = 0.5):#对掩体进行赋予颜色
    for c in range(3):
        image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - color_dark) + color_dark * color[c], image[:, :, c])
    return image


    

image_origin = cv2.imread('test_img/20240108094728.jpg')
image = cv2.cvtColor(image_origin, cv2.COLOR_BGR2RGB)
import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"

device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(image)

print(len(masks))
print(masks[0].keys())


image_select = image_origin.copy()
for i in range(len(masks)):
    color = tuple(np.random.randint(0, 256, 3).tolist())#随机列表颜色，就是
    selected_mask=masks[i]['segmentation']
    selected_image = apply_color_mask(image_select,selected_mask, color)
cv2.imwrite("res.jpg", selected_image)