当前位置:首页> AI教程> Segment Anything模型安装教程

Segment Anything模型安装教程

释放双眼,带上耳机,听听看~!
学习如何安装Segment Anything模型,包括环境准备、模型下载和测试使用。文章涵盖Python、PyTorch等内容。

这次给大家带来分割一切Segment Anything的安装教程,一起来看看吧~

Segment Anything模型安装教程

前期准备

Sam(Segment Anything Model)模型要求:

  • python>=3.8
  • pytorch>=1.7
  • torchvision>=0.8

我们先把Segment Anything的源代码下载下来: github.com/facebookres…
对git不熟的朋友可以下载zip然后解压:

Segment Anything模型安装教程

安装环境

查看python版本

python -V

当前环境不符合则
conda创建新环境:

conda create -n sam python=3.8

查看cuda版本

nvcc -V

根据cuda版本安装合适的pytorch:

网页可以ctrl+F搜索一下:cudatoolkit=11.1

https://pytorch.org/get-started/previous-versions/

我的是11.1,所以安装命令:

# CUDA 11.1
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge

进入项目根目录,安装sam

pip install -e .

安装其他依赖,sam的案例需要用到的相关库,因此需要装一下

pip install opencv-python pycocotools matplotlib onnxruntime onnx

下载模型

在上一节,我们安装好了环境和项目,但要想使用还需要下载一下模型,模型由于过大,所以一般都会和项目分开,sam一共有一下三种模型: Segment Anything模型安装教程 我们就下载默认的模型吧:

模型链接:github.com/facebookres… 下载好后,可以新建一个models文件夹放入

测试使用

导入依赖

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

展示结果函数,根据分区设置不同颜色

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((13)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

读取图片,案例图片在项目的notebooks/images里

image = cv2.imread('images/dog.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

加载模型,这里我用CPU跑,因为笔记本的3060显存只有6g跑不动,家人们,谁懂啊!

Segment Anything模型安装教程

如果你想用GPU,去掉注释行就行

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

sam_checkpoint = "../models/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)
plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show() 

结果:

Segment Anything模型安装教程

本期教程就结束啦,这里是GIS宇宙,我们下期再见~

商业合作请私信公众号
本人其他平台账号:

  • 微信公众号:GIS宇宙
  • CSDN:GIS_宇宙
  • 知乎:GIS宇宙
  • 掘金:GIS宇宙
本网站的内容主要来自互联网上的各种资源,仅供参考和信息分享之用,不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益,请联系我们,我们将尽快采取行动,包括删除或更正。
AI教程

专用AI加速器:能效、性能与模型复杂度

2023-12-9 11:35:14

AI教程

GPTCache:提高语言模型应用性能的利器

2023-12-9 11:49:14

个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索