SAM2 (Segment Anything 2) 돌려 보기

메타가 SAM 2를 공개한 지  2개월이 다 되어 갑니다. 

앞으로 이를 이용해 장난감 같은 기능 구현을 해 볼 생각인데 그때 객체 분할을 해야 할지도 몰라서 공개된 코드를 이용해 한번 돌려봤습니다. (참조: SAM 2 공식 깃헙)

 

절차는 다음과 같습니다. 

 

1.아나콘다 가상환경 생성

# 가상환경 생성
$ conda create -n sam2 python=3.11

 

2.모델 다운로드 및 필요 패키지 설치

# git clone 및 필요 패키지 설치 
$ git clone https://github.com/facebookresearch/segment-anything-2
$ cd segment-anything-2
$ pip3 install -e .

 

3.사전학습 가중치 다운로드

# weight 다운로드
$ cd checkpoints
$ ./download_ckpts.sh

 

4.테스트 이미지 준비

# 테스트 이미지 준비 
$ cp notebook/images/cars.jpg image.jpeg

 

 

5.CUDA 지원 환경 확인

코드 테스트를 하는 중에 오류가 발생 했습니다. 찾아보니 설치된 CUDA 버전이 맞지 않아 생긴 것이라고 하는데 제 경우, CUDA 12.4 버전이 적용되어 있었습니다. github의 Issue 부분에 해결 방법이 있네요. (링크)

 

$ vi /home/sol/git/segment-anything-2/sam2/modeling/sam/transformer.py

 

파일 상에 아래 설명처럼 Flash Attention 등의 설정 부분을 추가합니다. 

warnings.simplefilter(action="ignore", category=FutureWarning)
OLD_GPU, USE_FLASH_ATTN, MATH_KERNEL_ON = get_sdpa_settings()

# 아래 부분 추가
USE_FLASH_ATTN = False
MATH_KERNEL_ON = True
OLD_GPU = True

 

6.코드 테스트 

샘플 코드는 다음과 같습니다. 소스에 있는 테스트 이미지 가운데 다수의 자동차가 포함된 이미지로 테스트 합니다. 

import torch
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
from PIL import Image
import numpy as np

sam2_checkpoint = "checkpoints/sam2_hiera_large.pt"
model_cfg = "sam2_hiera_l.yaml"

sam2 = build_sam2(model_cfg, sam2_checkpoint, device ='cuda', apply_postprocessing=False)

mask_generator = SAM2AutomaticMaskGenerator(sam2)

image = np.array(Image.open('notebook/images/cars.jpg').convert("RGB"))
masks = mask_generator.generate(image)

 

7.결과 확인

객체 분할 테스트 결과는 다음과 같습니다. 생각보다 잘 나오는 것 같습니다. 

사진 : SAM 2를 이용해 이미지 상에서의 객체 분할한 결과 예시