SAM2をローカルGPU環境で動かす手順 (FastAPI + インタラクティブUI)

[project]
name = "segment2"
version = "0.1.0"
description = "Add your description here"
dependencies = [
    "huggingface-hub>=0.29.3",
    "torch>=2.6.0",
    "image>=1.5.33",
    "sam2>=1.1.0",
    "matplotlib>=3.10.1",
    "opencv-contrib-python>=4.11.0.86",
    "fastapi>=0.115.12",
    "uvicorn[standard]>=0.34.0",
    "python-multipart>=0.0.20",
    "opencv-python-headless>=4.11.0.86",
]
...

# 使用するモデルを選択 (tinyの方が精度は悪いがサイズは小さい)
SAM2_MODEL_NAME = "facebook/sam2-hiera-tiny"
# SAM2_MODEL_NAME = "facebook/sam2-hiera-large"

class Sam:
    def __init__(self):
        """SAM2 Predictor のインスタンスを作成する (FastAPI起動時)"""
        device = torch.device("cuda")
        if device.type == "cuda":
            torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
            # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
            if torch.cuda.get_device_properties(0).major >= 8:
                torch.backends.cuda.matmul.allow_tf32 = True
                torch.backends.cudnn.allow_tf32 = True
        else:
            raise NotImplementedError
        try:
            self.predictor = SAM2ImagePredictor.from_pretrained(SAM2_MODEL_NAME)
        except:
            self.predictor = None

    def predict_mask(
        self, image_path: str, input_points_list: list, input_labels_list: list
    ):
        # print(input_points_list)
        # print(input_labels_list)
        """指定された画像と点を使ってマスクを予測する (元のサンプルに近づける)"""
        if self.predictor is None:
            # load_model が FastAPI 起動時に呼ばれているはずだが、念のためチェック
            raise RuntimeError(
                "SAM2 predictor is not initialized. Application might not have started correctly."
            )

        try:
            img = Image.open(image_path).convert("RGB")
            img_array = np.array(img)
            image_np_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

            input_point_np = np.array(input_points_list, dtype=np.float32)
            input_label_np = np.array(input_labels_list, dtype=np.int32)

            # --- 元のサンプルのデバイス設定、autocast/TF32設定を推論前に適用 ---

            with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
                self.predictor.set_image(img_array)
                masks, scores, logits = self.predictor.predict(
                    point_coords=input_point_np,
                    point_labels=input_label_np,
                    multimask_output=True,
                )
                sorted_ind = np.argsort(scores)[::-1]
                masks = masks[sorted_ind]
                scores = scores[sorted_ind]
                logits = logits[sorted_ind]

                if len(scores) == 0:
                    print("Warning: No masks predicted.")
                    # マスクが見つからなかった場合、元画像を返す
                    return image_np_bgr, 0.0

                # 最もスコアの高いマスクを選択
                best_mask_idx = np.argmax(scores)
                best_mask_np = masks[best_mask_idx]
                best_score = scores[best_mask_idx]

                # マスクを元画像(BGR)に適用
                result_image_np = self.apply_mask_to_image(
                    image_np_bgr, best_mask_np, borders=True
                )

                return result_image_np, float(best_score)

        except FileNotFoundError:
            print(f"Error: Image file not found at {image_path}")
            return None, None
        except Exception as e:
            print(f"Error during prediction in sam_engine: {e}")
            import traceback

            traceback.print_exc()
            return None, None

    def apply_mask_to_image(
        self, image_np_bgr, mask_np, random_color=False, borders=True
    ):
        """マスクを画像に重ね合わせる (OpenCVを使用 - この部分は変更なし)"""
        if mask_np.dtype != bool:
            mask_boolean = mask_np.astype(bool)
        else:
            mask_boolean = mask_np

        if random_color:
            color = np.array(
                [
                    np.random.randint(0, 256),
                    np.random.randint(0, 256),
                    np.random.randint(0, 256),
                ],
                dtype=np.uint8,
            )
        else:
            color = np.array([255, 144, 30], dtype=np.uint8)  # BGR: 青っぽい色

        h, w = mask_boolean.shape[-2:]
        mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
        mask_colored[mask_boolean] = color

        alpha = 0.5
        blended = cv2.addWeighted(image_np_bgr, 1.0, mask_colored, alpha, 0)

        if borders:
            try:
                mask_uint8 = mask_boolean.astype(np.uint8) * 255
                contours, _ = cv2.findContours(
                    mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
                )
                cv2.drawContours(blended, contours, -1, (255, 255, 255), thickness=2)
            except Exception as e:
                print(f"Warning: Failed to draw contours: {e}")

        return blended

mkdir frontend
touch frontend/index.html

// --- 主要なDOM要素を取得 ---
const imageUpload = document.getElementById('image-upload');
const imageContainer = document.getElementById('image-container');
const uploadedImage = document.getElementById('uploaded-image');
const segmentBtn = document.getElementById('segment-btn');
const resultImage = document.getElementById('result-image');
const resultScore = document.getElementById('result-score');
// (他の要素取得は省略)

// --- 状態変数 ---
let currentFilename = null;     // サーバー上の画像ファイル名
let imageOriginalWidth = 0;     // 元画像の幅
let imageOriginalHeight = 0;    // 元画像の高さ
let points = [];                // クリックされた点の情報 [{originalX, originalY, label}, ...]


// --- 1. 画像アップロード処理 (抜粋) ---
imageUpload.addEventListener('change', async (event) => {
    const file = event.target.files[0];
    if (!file) return;
    // (UIリセット処理などは省略)

    const formData = new FormData();
    formData.append('file', file);

    try {
        // FastAPIの /upload/ エンドポイントに画像をPOST
        const response = await fetch('/upload/', { method: 'POST', body: formData });
        if (!response.ok) throw new Error('Upload failed');

        const data = await response.json();
        currentFilename = data.filename;      // ファイル名を取得
        imageOriginalWidth = data.width;    // 元画像のサイズを取得
        imageOriginalHeight = data.height;

        // 画像プレビュー表示 (FileReaderを使った処理は省略)
        // ... uploadedImage.src = ...
        console.log(`Uploaded: ${currentFilename}, Size: ${imageOriginalWidth}x${imageOriginalHeight}`);

    } catch (error) {
        console.error('Upload error:', error);
        // (エラー表示処理は省略)
    }
});

// --- 2. 画像クリック時の座標取得・保存 (抜粋) ---
imageContainer.addEventListener('click', (event) => {
    if (!currentFilename || !uploadedImage.complete) return; // 画像ロード完了前にクリックされるのを防ぐ

    const rect = imageContainer.getBoundingClientRect();
    const displayWidth = uploadedImage.offsetWidth;
    const displayHeight = uploadedImage.offsetHeight;

    // クリック座標 (表示されている画像に対する相対座標)
    const clickX = event.clientX - rect.left;
    const clickY = event.clientY - rect.top;

    // 元画像における座標にスケーリング
    const scaleX = imageOriginalWidth / displayWidth;
    const scaleY = imageOriginalHeight / displayHeight;
    const originalX = Math.round(clickX * scaleX);
    const originalY = Math.round(clickY * scaleY);

    // 画像範囲外なら無視
    if (originalX < 0 || originalX >= imageOriginalWidth || originalY < 0 || originalY >= imageOriginalHeight) return;

    // 選択されているラベル (1: Positive, 0: Negative) を取得
    const selectedLabel = parseInt(document.querySelector('input[name="click-mode"]:checked').value);

    // 点の情報を保存 (UIに点を描画する処理は省略)
    points.push({ originalX, originalY, label: selectedLabel });
    console.log('Point added:', { originalX, originalY, label: selectedLabel });
    segmentBtn.disabled = false; // 実行ボタンを有効化
});

// --- 3. セグメンテーション実行 (抜粋) ---
segmentBtn.addEventListener('click', async () => {
    if (!currentFilename || points.length === 0) return;
    // (ローディング表示などのUI制御は省略)

    // バックエンドAPIに送るデータを作成
    const requestData = {
        filename: currentFilename,
        points: points.map(p => [p.originalX, p.originalY]), // [[x1, y1], [x2, y2], ...]
        labels: points.map(p => p.label),                   // [1, 0, ...]
    };

    try {
        // FastAPIの /segment/ エンドポイントにJSONデータをPOST
        const response = await fetch('/segment/', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(requestData),
        });
        if (!response.ok) throw new Error('Segmentation failed');

        const data = await response.json();
        // 結果画像(Base64)を表示
        resultImage.src = data.result_image;
        resultImage.classList.remove('hidden');
        // スコア表示
        resultScore.textContent = data.score ? `予測スコア: ${data.score.toFixed(3)}` : 'スコアなし';

    } catch (error) {
        console.error('Segmentation error:', error);
        // (エラー表示処理は省略)
    } finally {
        // (ローディング解除などのUI制御は省略)
    }
});

// --- 4. 点をクリアする処理 (抜粋) ---
function clearPoints() {
    points = []; // 配列を空にする
    // (UI上のマーカー削除処理は省略)
    console.log('Points cleared');
    segmentBtn.disabled = true;
}

// clearPointsBtn のクリックイベントリスナーで clearPoints() を呼び出す (コード省略)

rye run uvicorn main:app --host 0.0.0.0 --port 8000