다음으로 SAM 모델을 로드하고 초기화합니다. 이를 통해 모델을 사용할 준비를 마치고, 이후에는 이미지 데이터를 불러와 모델을 통해 분석하거나 변형하는 작업을 진행할 수 있습니다.

import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor # ①

CHECKPOINT_PATH = "sam_vit_h_4b8939.pth"
MODEL_TYPE = "vit_h"
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device='cuda:0')

sam_model_registry는 사용 가능한 SAM 모델들을 담고 있는 레지스트리, SamAutomatic MaskGenerator는 자동 마스크 생성기, SamPredictor는 SAM 모델의 예측을 수행하는 클래스입니다(①). 변수들을 선언한 후 sam_model_registry에서 지정한 타입의 모델을 로드하고, 사전 훈련된 가중치를 이용해 초기화합니다. .to(device=DEVICE)는 모델을 특정 디바이스(예 CPU 또는 cuda:0)에 로드하기 위한 코드입니다. 여기서 DEVICE 변수는 실행 환경에 따라 CPU 또는 GPU 등으로 설정되어야 합니다.