Segment Anything

# pip install segment-anything opencv-python matplotlib torch torchvision
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamPredictor
import cv2

# 모델 로드 (vit_h, vit_l, vit_b 중 선택)
sam_checkpoint = "sam_vit_h_4b8939.pth"  # 다운로드 필요
model_type = "vit_h"

device = "cuda"  # GPU 사용, CPU는 "cpu"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

predictor = SamPredictor(sam)

# 이미지 로드
image = cv2.imread("your_image.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 이미지 임베딩 생성 (한 번만 실행)
predictor.set_image(image)

# 점 프롬프트로 세그멘테이션
# (x, y) 좌표, label 1=전경, 0=배경
input_point = np.array([[500, 375]])
input_label = np.array([1])

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,  # 여러 마스크 후보 반환
)

# 가장 높은 점수의 마스크 선택
best_mask_idx = np.argmax(scores)
best_mask = masks[best_mask_idx]

# 결과 시각화
def show_mask(mask, ax, color=(30/255, 144/255, 255/255, 0.6)):
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * np.array([*color[:3], color[3]])
    ax.imshow(mask_image)

fig, ax = plt.subplots(1, 1, figsize=(10, 10))
ax.imshow(image)
show_mask(best_mask, ax)
ax.scatter(input_point[:, 0], input_point[:, 1], c='red', s=200, marker='*')
ax.set_title(f"Score: {scores[best_mask_idx]:.3f}")
ax.axis('off')
plt.savefig("segmentation_result.png", bbox_inches='tight')
plt.show()

print(f"분할된 마스크 개수: {len(masks)}")
print(f"각 마스크 점수: {scores}")

# pip install sam2
import torch
from sam2.build_sam import build_sam2_video_predictor

# SAM 2 비디오 모델 로드
predictor = build_sam2_video_predictor(
    "sam2_hiera_large.yaml",
    "sam2_hiera_large.pt"
)

# 비디오 초기화
video_path = "your_video.mp4"
inference_state = predictor.init_state(video_path=video_path)

# 첫 프레임에서 객체 선택 (점 프롬프트)
frame_idx = 0
obj_id = 1
points = np.array([[500, 375]], dtype=np.float32)
labels = np.array([1], dtype=np.int32)

# 프롬프트 추가
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=frame_idx,
    obj_id=obj_id,
    points=points,
    labels=labels,
)

# 전체 비디오 전파
video_segments = {}
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(
    inference_state
):
    video_segments[out_frame_idx] = {
        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
        for i, out_obj_id in enumerate(out_obj_ids)
    }

print(f"총 {len(video_segments)} 프레임 세그멘테이션 완료")