[Split everything, everyone loses their jobs together! ]Segment Anything lightspeed installation tutorial

This time, I will bring you the installation tutorial of Segment Anything, let’s take a look~

Preparation

Sam (Segment Anything Model) model requirements:

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

Let's download the source code of Segment Anything first:
https://github.com/facebookresearch/segment-anything

Friends who are not familiar with git can download the zip and unzip it:

Installation Environment

Check python version

python -V

If the current environment does not match,

conda creates a new environment:

conda create -n sam python=3.8

Check cuda version

nvcc -V

Install the appropriate pytorch according to the cuda version:

The webpage can be searched by ctrl+F: cudatoolkit=11.1

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

Mine is 11.1, so the installation command:

# 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

Enter the project root directory and install sam

pip install -e .

Install other dependencies, related libraries needed in the case of sam, so you need to install them

pip install opencv-python pycocotools matplotlib onnxruntime onnx

download model

In the previous section, we installed the environment and project, but we need to download the model to use it. Because the model is too large, it is usually separated from the project. There are three models in sam: let’s download the default model

. :

Model link: https://github.com/facebookresearch/segment-anything#model-checkpoints
After downloading, you can create a new models folder and put it in

test use

import dependencies

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

Display the result function, set different colors according to the partition

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((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

Read the picture, the case picture is in the notebooks/images of the project

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

Load the model, here I use the CPU to run, because the 3060 video memory of the notebook is only 6G, and the family members, who knows!

If you want to use the GPU, just uncomment the line

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)

Get the recognized area and draw it

masks = mask_generator.generate(image)
plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

result:

This tutorial is over, here is the GIS universe, see you next time~

For business cooperation, please private message the

official account My other platform account:

WeChat public account: GIS universe
CSDN: GIS_Universe
Zhihu: GIS Universe
Nuggets: The GIS Universe