on
딥러닝 - SAM (Segment Anything Model)을 이용한 이미지 마스킹
딥러닝 - 지난 글
목차
SAM (Segment Anything Model)이란?
-
SAM은 입력한 텍스트 프롬프트나 클릭한 지점을 기반으로 AI가 이미지 내에서 특정 물체를 분리해주는 이미지 분할 모델이다. 이미지 내에서 특정 물체를 식별하고 분리할 수 있게 해준다.
-
“이미지 세그멘테이션 (segmentation)”을 위한 모델이며, 픽셀이 어떤 객체에 속해있는지를 식별하는 알고리즘이다.
-
Meta는 SAM 과 40억개의 파라미터 데이터셋 (SA-1B) 모델을 공개하였다.
-
SAM은 물체가 무엇인지에 대한 일반적인 개념을 학습 했고, 훈련중에 만나지 못한 물체 및 이미지 유형에 대해서도 이미지/비디오의 모든 객체에 대해서 마스크를 생성 가능 추가적인 훈련 없이도 새로운 이미지 도메인(물속 사진이나, 세포 현미경 사진등)에도 적용 가능하다.
SAM의 목표
다음은 Meta에서 공개한 SAM에 대한 소개 글을 요약한 것이다. Segment Anything 프로젝트는 무엇이며 목표는 무엇인가요?
Segment Anything 프로젝트는 이미지 세분화를 위한 새로운 작업, 데이터 세트, 모델을 도입하여 이미지 세분화를 대중화하는 것을 목표로 합니다. 이 프로젝트의 목표는 이미지 세분화를 위한 작업별 모델링 전문 지식, 학습 컴퓨팅, 맞춤형 데이터 주석의 필요성을 줄이는 것입니다. 이 프로젝트에서는 Segment Anything Model(SAM)과 Segment Anything 10억 마스크 데이터세트(SA-1B)를 도입합니다. SAM은 다양한 데이터에 대해 학습되어 특정 작업에 적응할 수 있는 즉시 사용 가능한 모델입니다. SA-1B는 연구 목적으로 공개된 역대 최대 규모의 세분화 데이터 세트입니다. 이 프로젝트는 광범위한 애플리케이션을 지원하고 컴퓨터 비전을 위한 기초 모델에 대한 추가 연구를 촉진하는 것을 목표로 합니다. 자연어 처리 모델의 프롬프트와 유사한 이미지 세분화를 위한 기초 모델을 만드는 것이 목표입니다. SAM 모델은 광범위한 사용 사례를 포괄할 수 있을 만큼 일반적이며 추가 교육 없이도 새로운 이미지 ‘도메인’에 즉시 사용할 수 있습니다. 궁극적인 목표는 이미지에서 물체를 찾고 세그먼트화해야 하는 다양한 영역의 애플리케이션을 강화하는 것입니다.
SAM의 원리
SAM의 원리는 다음 참고 자료에 잘 설명 되어 있다.
Segment Anything 논문 정리
박진우 GitHub
SAM GitHub
GitHub: https://github.com/facebookresearch/segment-anything
Demo: https://segment-anything.com/demo
SAM 테스트를 위한 준비
로컬에서 SAM 테스트를 위해서는 다른 딥러닝 알고리즘과 마찬가지로 약간의 준비가 필요하다. 본인은 Anaconda 가상환경에서 테스트하였으며, Anaconda 가상환경 세팅 방법 부터 Pytorch와 Cuda 설치 및 버전 컨트롤까지 다루어 보도록 하겠다.
Anaconda 가상환경 생성
Anaconda 가상환경을 만드는 방법은 다음과 같다.
SAM의 경우 python>=3.8이상의 버전을 요구한다. 따라서 본인은 3.9버전을 통해 테스트를 진행하였다.
Windows의 경우 Anaconda Prompt, Mac의 경우 터미널에서 진행한다.
conda create -n 가상환경이름 python=3.9
여기서 y를 타이핑하고 생성을 진행한다.
가상환경이 잘 생성되었는 지 확인해보자.
conda env list
다음처럼 생성되었으면 성공이다.
생성된 가상환경을 적용해보자.
conda activate 가상환경 이름
다음처럼 (base)가 (내가 만든 가상환경 이름)으로 바뀌었으면 성공이다.
SAM 설치하기
SAM의 git 주소를 클론하여 설치하기 위해서는 git을 설치해야 한다.
conda install git pip
SAM의 깃허브 주소를 통해 클론 설치를 할 수 있다.
git clone git+https://github.com/facebookresearch/segment-anything.git
해당 방법이 불가능하다면, 수동으로 직접 다운로드하여 설치할 수 있다.
다음과 같이 DownloadZIP을 클릭하여 압축파일을 다운받고 원하는 폴더로 이동시키면 수동 설치가 완료된다.
이어서 SAM의 Pretrained model을 설치해야 한다.
다음에서 default or vit_h: ViT-H SAM model. 을 클릭하여
설치한 후 SAM 설치 폴더에 model폴더를 만들어 넣어주면 된다. 용량은 총 2.4GB이다.
다음으로 CUDA 설치를 진행할 것이다. CUDA는 NVIDIA 그래픽카드가 설치된 로컬에만 설치 및 이용이 가능하다. 따라서 해당 그래픽카드가 설치되어있지 않은 경우 바로 다음 Pytorch 설치 과정으로 넘어가면 된다.
본인 로컬 환경에 맞는 CUDA 버전은 다음 글을 참고하여 확인할 수 있다.
CUDA 설치 코드는 다음사이트에서 확인할 수 있다.
본인은 CUDA 11.6 버전에 해당하기 때문에 다음과 같이 설치할 것이다.
conda install -c "nvidia/label/cuda-11.6.2" cuda-toolkit
다음 방법으로 설치된 CUDA 버전을 확인할 수 있다.
nvcc --version
다음과 같이 11.6이 확인되면 올바르게 설치된 것이다.
다음으로 Pytorch를 설치할 것이다. SAM은 pytorch>=1.7 과 torchvision>=0.8 버전을 지원하기 때문에 해당 버전에 맞게 설치를 진행할 것이다.
Pytorch는 최근 2.0으로 업데이트 되었다. 따라서 홈페이지에 접속시 Pytorch 2.0을 기본으로 설치하게 되어있을 것이다. 하지만 본인은 이전 버전을 이용하여 테스트를
진행할 것이다. 이유는 최신 버전은 대체로 하위버전도 지원하지만 불안정하거나 실행시 오류가 발생하는 경우가 있기 때문에 안전하게 하위 버전으로 진행하는 것이 좋다는
개인적인 경험에 의한 것이기 때문이다…
Pytorch 이전 버전은 다음 사이트에서 설치 가능하다. 본인은 v1.13.1 버전을 설치 및 테스트 하였다.
만약 그래픽 카드 없이 CPU로만 테스트 할 경우 다음 명령어를 입력하여 설치할 수 있다.
# CPU Only
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 cpuonly -c pytorch
CUDA를 설치하고 온 경우에 본인의 CUDA 버전에 맞게 설치할 수 있다. 본인은 11.6이기 때문에 해당 버전의 Pytorch를 설치하겠다.
# CUDA 11.6
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
이제 SAM 테스트를 위한 모든 준비가 끝났다. 이어서 SAM 테스트를 진행해보자.
SAM 테스트
notebooks 폴더에는 automatic_mask_generator_example.ipynb, predictor_example.ipynb 두 개의 파일이 존재한다. 본인은 이 두 개의 파일을
테스트해 보았다.
mask_generator의 경우 해당 이미지에 있는 모든 객체를 segment 해준다. predictor의 경우 해당 이미지에 포인터를 적용하고, 포인터가 있는 지점의 객체를
segment한다.
다음 토마토 이미지를 이용하여 SAM을 테스트 해보았다.
mask_generator
import numpy as np
import torch
import gc
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((1, 3)).tolist()[0]
for i in range(3):
img[:,:,i] = color_mask[i]
ax.imshow(np.dstack((img, m*0.35)))
image = cv2.imread('images/tomato.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(20,20))
plt.imshow(image)
plt.axis('off')
plt.show()
import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
sam_checkpoint = "../model/sam_vit_h_4b8939.pth"
model_type = "vit_h"
### CPU 사용자의 경우
# device = 'cpu'
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())
output
114
dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])
모델 결과 확인
plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()
output
predictor
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
def show_mask(mask, ax, random_color=False):
if random_color:
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
else:
color = np.array([30/255, 144/255, 255/255, 0.6])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
ax.imshow(mask_image)
def show_points(coords, labels, ax, marker_size=375):
pos_points = coords[labels==1]
neg_points = coords[labels==0]
ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
def show_box(box, ax):
x0, y0 = box[0], box[1]
w, h = box[2] - box[0], box[3] - box[1]
ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))
image = cv2.imread('images/tomato.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(10,10))
plt.imshow(image)
plt.axis('on')
plt.show()
import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor
sam_checkpoint = "../model/sam_vit_h_4b8939.pth"
model_type = "vit_h"
### CPU 사용자의 경우
# device = 'cpu'
device = "cuda"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
predictor.set_image(image)
# 포인터 설정
input_point = np.array([[400, 400]])
input_label = np.array([1])
plt.figure(figsize=(10,10))
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
plt.axis('on')
plt.show()
output
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=True,)
masks.shape # (number_of_masks) x H x W
for i, (mask, score) in enumerate(zip(masks, scores)):
plt.figure(figsize=(10,10))
plt.imshow(image)
show_mask(mask, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
plt.axis('off')
plt.show()
output
결론
대단하다…