[SAM] 모델 다운로드 및 실행

Segment Anything Model (SAM, Meta AI)

• 미리 학습된 강력한 세그멘테이션 모델
• 사람이 따로 학습 데이터를 만들 필요 없이 미술 작품을 자동으로 분할 가능

 

모델 다운로드

https://github.com/facebookresearch/segment-anything?tab=readme-ov-file

GitHub - facebookresearch/segment-anything: The repository provides code for running inference with the SegmentAnything Model (S

 

아래로 내리다 보면, Model Checkpoints에 여러 모델이 있다.

VIT-H SAM model 클릭하여 다운로드.

python 코드와 같은 위치에, 다운로드한 sam_vit_h.pth 파일을 복붙하자.

 

 

 

from segment_anything import SamPredictor, sam_model_registry

# ----------------------------- SAM 설정 -----------------------------
sam = sam_model_registry['vit_h'](checkpoint='app/models/one_imageDetection/sam_vit_h.pth')
sam_predictor = SamPredictor(sam)

# image 넣어줘야함
def segment_with_sam(image, bbox):
    sam_predictor.set_image(image)
    bbox = np.array(bbox).reshape(1, 4)
    masks, _, _ = sam_predictor.predict(box=bbox)
    return masks[0]