Segment Anything 사용해 보기

KYYLE·2024년 4월 13일

Tutorials

목록 보기

6/9

Segment Anything

Segment Anything (SAM) 모델은 Meta에서 발표한 segmentation task를 위한 foundation model입니다. Foundation model 답게, 다양한 상황 및 이미지에서 훌륭한 segmentation 결과를 보여줍니다.

SAM을 사용해 보려면 Meta에서 제공하는 웹 데모 사이트를 사용하거나, 공식 깃허브에서 모델 체크포인트(가중치)를 내려받아 사용할 수 있습니다.

공식 깃허브에서는 모델 체크포인트와 함께 간단한 튜토리얼(notebooks/ 디렉터리 아래 *.ipynb)을 제공해 주는데요, 제공되는 튜토리얼처럼 prompt 좌표를 미리 설정해서 입력하는 것이 아닌 사용자가 직접 마우스로 클릭하여 prompt를 넣을 수 있도록 간단히 파이썬 코드를 작성해 보았습니다.

Installation

먼저, Meta의 공식 깃허브에서 모델 체크포인트를 내려받습니다. 저는 ViT-H SAM model 버전을 사용하였습니다. 다른 버전을 사용해도 되지만, 다른 체크포인트를 사용한다면 이후 코드의 일부분을 수정해야 합니다.

다음으로, 제 깃허브의 SAM-Inference 리포지터리를 클론한 다음 다운로드한 체크포인트 파일을 폴더 안에 넣습니다.

프로그램을 실행하기 위해서는 최소 아래의 라이브러리가 필요합니다.

opencv-python, numpy, torch, matplotlib
GPU에 맞는 pytorch와 CUDA 사용을 권장해 드립니다.
추가로, opencv의 cv2.imshow() 함수를 실행할 수 있어야 합니다. 일부 리눅스 환경에서는 지원되지 않을 수 있습니다.

How to Use

SAM-Inference 디렉터리에서 터미널을 실행한 후, python sam_inference.py를 실행합니다. 만약 다른 체크포인트를 사용하신다면, sam_inference.py 파일 내부를 다음과 같이 변경합니다.

# setting 
sam_checkpoint = "sam_vit_h_4b8939.pth" # 다운로드 받은 체크포인트 파일 
model_type = "vit_h" # vit_l 또는 vit_b
device = "cuda:0"

프로그램을 실행하면 먼저 이미지의 경로를 입력해야 합니다. 깃허브를 클론하면 기본적으로 imgs/bird.jpg, imgs/car.jpg 파일이 존재합니다. 해당 이미지를 사용하거나, 원하는 이미지를 사용할 수 있습니다.

이미지 경로를 올바르게 입력했다면, 다음으로 segmentation mode를 설정합니다. 지원되는 mode는 두 가지로, Single object segmentation mode와 Multiple object segmentation mode입니다. 말 그대로 하나의 객체만 segmentation 하거나, 여러 개의 객체를 segmentation 하는 mode입니다.

Single object segmentation mode에서는 여러 개의 point prompt와 하나의 box prompt를 지원합니다. Point prompt의 label은 양성이거나 음성일 수 있습니다. Multiple object segmentation mode에서는 여러 개의 box prompt를 지원합니다.

Segmentation mode를 설정하면, 다음으로 'p' 또는 'b'를 눌러 point 또는 box를 이미지 위에 그릴 수 있습니다(모두 소문자입니다).

'p'를 누르면 Point mode로 진입하여, 마우스 좌클릭 혹은 우클릭으로 이미지 위에 point를 찍을 수 있습니다. 좌클릭은 초록색 점이 생기며, 이는 positive label을 의미합니다. 우클릭은 빨간색 점이 생기며, 이는 negative label을 의미합니다.

'b'를 누르면 Bbox mode로 진입하여, 마우스 드래그를 통해 이미지 위에 box를 그릴 수 있습니다.

'q'를 누르면 종료, 혹은 뒤로 가기 기능을 수행합니다. 'p'를 눌러 point prompt를 생성한 후, 'q'를 눌러 Point mode를 종료한 다음 'b'를 눌러 box prompt를 생성할 수 있습니다.

잘못된 위치에 prompt를 생성했거나 새로운 결과를 얻고 싶은 경우, 'r'을 눌러 생성된 모든 prompt를 제거할 수 있습니다.

이미지 위에 prompt를 모두 생성하였다면, 'i'를 눌러 모델의 출력 결과를 확인할 수 있습니다. 새로운 window에서 segmentation 결과가 출력되며, result/ 디렉터리 아래에 segmentation 결과가 저장됩니다.