from segment_anything import sam_model_registry, SamPredictor
import cv2

# 모델 로드
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
sam = sam.to('cuda')
predictor = SamPredictor(sam)

# 이미지 설정
image = cv2.imread("image.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
predictor.set_image(image)

import numpy as np

# 클릭 좌표로 세그멘테이션 (포그라운드 점)
input_point = np.array([[500, 375]])
input_label = np.array([1])  # 1: 포그라운드, 0: 배경

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,  # 3개의 마스크 후보 생성
)

# 가장 높은 점수의 마스크 선택
best_mask = masks[scores.argmax()]
print(f"마스크 형태: {best_mask.shape}")  # [H, W] boolean
print(f"최고 점수: {scores.max():.4f}")

# 바운딩 박스로 세그멘테이션
input_box = np.array([100, 150, 400, 450])  # [x1, y1, x2, y2]

masks, scores, logits = predictor.predict(
    box=input_box,
    multimask_output=False,
)

from segment_anything import SamAutomaticMaskGenerator

# 자동 마스크 생성기
mask_generator = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,       # 그리드 밀도
    pred_iou_thresh=0.86,     # IoU 임계값
    stability_score_thresh=0.92,
    min_mask_region_area=100,  # 최소 마스크 면적
)

masks = mask_generator.generate(image)
print(f"생성된 마스크 수: {len(masks)}")

# 마스크 정렬 (면적 순)
sorted_masks = sorted(masks, key=lambda x: x['area'], reverse=True)
for i, mask in enumerate(sorted_masks[:5]):
    print(f"마스크 {i}: 면적={mask['area']}, IoU={mask['predicted_iou']:.3f}")

import torch
from sam2.build_sam import build_sam2_video_predictor

# SAM2 비디오 예측기
predictor = build_sam2_video_predictor(
    "sam2_hiera_large.yaml",
    "sam2_hiera_large.pt",
)

# 비디오 초기화
with torch.inference_mode():
    state = predictor.init_state(video_path="video.mp4")

    # 첫 프레임에서 프롬프트 (클릭)
    _, _, masks = predictor.add_new_points_or_box(
        state, frame_idx=0,
        obj_id=1,
        points=np.array([[400, 300]]),
        labels=np.array([1]),
    )

    # 전체 비디오 전파
    for frame_idx, obj_ids, masks in predictor.propagate_in_video(state):
        print(f"프레임 {frame_idx}: {len(obj_ids)}개 객체 추적 중")