본문 바로가기

인공지능

Segment Anything Model(SAM) 사용하기

Meta AI Research (Facebook)에서 제작한 SAM을 사용하는 방법에 대해 포스팅 하겠다.

https://bestkcs1234.tistory.com/73 에서 리뷰한 모델이고 페이스북에서 아주 대용량의 데이터를 이용해 제작한 모델이다.

 

우선 해당 깃허브 https://github.com/facebookresearch/segment-anything 에 접속해서 repo를 다운받아 준다.

1. 직접 알집파일로 다운받기
2. 파이참 터미널에서 git clone 하기
2-1. git clone https://github.com/facebookresearch/segment-anything.git    입력하여 깃허브를 현 directory에 clone
2-2. cd segment-anything; pip install -e .     입력하여 필요한 dependencies 다운

 

 

레포지토리를 다운 받은 다음 pretrained-weight를 이용하고 싶을 경우는 깃허브를 내리다보면 해당하는 vit_h / vit_l / vit_b 의 weight를 다운받는다.

필자는 vit_b를 다운받았으며 sam_vit_b_01ec64.pth 라는파일을 segment_anything 디렉터리로 옮겨준다

 

 

 

필자의 directory 구성은 왼쪽과 같다. SAM을 사용해보기 위한 SAM_test 스크립트를 하나 생성했다. (깃허브의 다른 파일들은 삭제)

또한 아래의 코드들은 파이참의 "Run"을 이용한 코드가 아닌 python console 환경에서 실행한 것이기 때문에 path와 같은 코드들은 본인의 환경에 맞도록 고치길 바란다.

 

from code_zip.segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "code_zip/segment_anything/sam_vit_b_01ec64.pth"
model_type = "vit_b"

device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

predictor = SamPredictor(sam)

깃허브에서 제시한 해당 코드를 model_type = 'vit_b', sam_checkpoint = 'sam_vit_b_01ec64.pth'로 변경해준다.

Root directory가 아닌 다른 directory에 저장하였더니 no modeul 에러가 뜬다.

이는 SamPredictor를 import 하면서 생긴 Path 오류이기 때문에 predictor.py의 코드를 고치던지 segment_anthing direcotry를 root directory로 옮겨주어야 한다.

 

path 오류를 해결하면 다음과 같이 모델이 로딩된다.

 

 

 

이 이후에는 해당 깃허브의 notebook 란들을 참고하면 된다.

1. object masking - (https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb)

 

segment-anything/notebooks/predictor_example.ipynb at main · facebookresearch/segment-anything

The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model. -...

github.com

 

 

2. automatic mask generator - (https://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb)

 

segment-anything/notebooks/automatic_mask_generator_example.ipynb at main · facebookresearch/segment-anything

The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model. -...

github.com