SAMによる画像セグメンテーション

REM Python をシステム領域にインストール
winget install --scope machine --id Python.Python.3.12 -e --silent
REM Python のパス設定
set "PYTHON_PATH=C:\Program Files\Python312"
set "PYTHON_SCRIPTS_PATH=C:\Program Files\Python312\Scripts"
echo "%PATH%" | find /i "%PYTHON_PATH%" >nul
if errorlevel 1 setx PATH "%PATH%;%PYTHON_PATH%" /M >nul
echo "%PATH%" | find /i "%PYTHON_SCRIPTS_PATH%" >nul
if errorlevel 1 setx PATH "%PATH%;%PYTHON_SCRIPTS_PATH%" /M >nul

winget install --scope machine Codeium.Windsurf -e --silent

# SAMリアルタイム画像セグメンテーション
# 特徴技術名: Segment Anything Model (SAM)
# 出典: A. Kirillov et al., "Segment Anything," in Proc. IEEE/CVF International Conference on Computer Vision (ICCV), 2023, pp. 4015-4026.
# 特徴機能: ゼロショット汎用セグメンテーション - 追加学習なしで任意の物体を自動検出・分割
# 学習済みモデル: SAM ViT-Base（91Mパラメータ、375MB、汎用物体セグメンテーション、https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth）
# 方式設計
#   - 関連利用技術:
#     * PyTorch（深層学習フレームワーク、GPU加速対応）
#     * OpenCV（画像処理、カメラ制御、表示機能）
#     * scikit-learn（特徴量正規化、MinMaxScaler）
#     * Pillow（日本語テキスト描画）
#   - 入力と出力: 入力: 動画（ユーザは「0:動画ファイル，1:カメラ，2:サンプル動画」のメニューで選択．0:動画ファイルの場合はtkinterでファイル選択．1の場合はOpenCVでカメラが開く．2の場合はhttps://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.aviを使用）、出力: セグメンテーション結果（OpenCV画面でリアルタイム表示、各検出物体を色分け表示、検出物体数を画面表示）
#   - 処理手順: 1.フレーム取得 2.BGR→RGB変換 3.SAM自動マスク生成 4.特徴量抽出（重心・アスペクト比） 5.特徴量正規化 6.色付けマスク生成 7.画像合成・テキスト描画・表示
#   - 前処理: BGR→RGB色空間変換（SAMの入力形式に適合）
#   - 後処理: 特徴量正規化による色分け、重み付き画像合成（元画像60%、マスク40%）、検出情報のオーバーレイ表示
#   - 調整を必要とする設定値: points_per_side（デフォルト32）: グリッドポイント数、pred_iou_thresh（デフォルト0.88）: マスク品質閾値
# 将来方策: pred_iou_thresh値の動的調整機能 - フレーム内の物体数に応じて閾値を自動調整し、最適な検出精度を実現
# その他の重要事項: CUDA対応、DirectShowバックエンド使用
# 前準備: pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
#         pip install opencv-python matplotlib numpy scikit-learn Pillow
#         pip install git+https://github.com/facebookresearch/segment-anything.git

import cv2
import numpy as np
import torch
import tkinter as tk
from tkinter import filedialog
import urllib.request
import time
from datetime import datetime
from pathlib import Path
from sklearn.preprocessing import MinMaxScaler
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
from PIL import Image, ImageDraw, ImageFont

# GPU/CPU自動選択
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'デバイス: {str(device)}')

# 設定定数
CHECKPOINT_NAME = 'sam_vit_b_01ec64.pth'
MODEL_TYPE = 'vit_b'
BASE_URL = 'https://dl.fbaipublicfiles.com/segment_anything/'
SEED = 42

# 調整可能な設定値
POINTS_PER_SIDE = 32  # グリッドポイント数
PRED_IOU_THRESH = 0.88  # マスク品質閾値

# 画像処理設定
ORIGINAL_WEIGHT = 0.6  # 元画像の重み
MASK_WEIGHT = 0.4      # マスクの重み
DEFAULT_COLOR = 0.5    # 単一マスク時のデフォルト色値

# フォント設定
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE = 30

# 乱数シード設定（再現性確保）
np.random.seed(SEED)
torch.manual_seed(SEED)

print('=== SAMリアルタイム画像セグメンテーション ===')
print('概要: Segment Anything Model (SAM)を使用した自動物体検出')
print('出力: 検出された物体を色分けして表示')
print('=====================================')

# SAMモデルダウンロード
if not Path(CHECKPOINT_NAME).exists():
    print(f'ダウンロード中: {CHECKPOINT_NAME}')
    urllib.request.urlretrieve(f'{BASE_URL}{CHECKPOINT_NAME}', CHECKPOINT_NAME)
    print(f'ダウンロード完了: {CHECKPOINT_NAME}')
else:
    print(f'既存ファイルを使用: {CHECKPOINT_NAME}')

# SAMモデル初期化
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_NAME)
sam.to(device=device)
mask_generator = SamAutomaticMaskGenerator(
    sam,
    points_per_side=POINTS_PER_SIDE,
    pred_iou_thresh=PRED_IOU_THRESH
)

frame_count = 0
results_log = []

def video_frame_processing(frame):
    global frame_count
    current_time = time.time()
    frame_count += 1

    # BGR to RGB変換
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # SAM自動マスク生成
    masks = mask_generator.generate(frame_rgb)

    # マスク画像初期化
    mask_img = np.zeros((frame_rgb.shape[0], frame_rgb.shape[1], 3), dtype=np.uint8)

    if masks:
        # 特徴量抽出処理
        features_list = []
        for mask_data in masks:
            mask = mask_data['segmentation']

            # 重心計算
            M = cv2.moments(mask.astype(np.uint8))
            if M['m00'] != 0:
                cx = int(M['m10'] / M['m00'])
                cy = int(M['m01'] / M['m00'])
            else:
                cx, cy = 0, 0

            # アスペクト比計算
            x, y, w, h = cv2.boundingRect(mask.astype(np.uint8))
            aspect_ratio = float(w) / h if h > 0 else 0

            features = [cx, cy, aspect_ratio]
            features_list.append(features)

        # 特徴量正規化処理
        feat_matrix = np.array(features_list)
        if len(feat_matrix) == 1:
            norm_feat = np.array([[DEFAULT_COLOR, DEFAULT_COLOR, DEFAULT_COLOR]])
        else:
            scaler = MinMaxScaler()
            norm_feat = scaler.fit_transform(feat_matrix)

        # 色付けマスク生成処理
        for i, mask_data in enumerate(masks):
            mask = mask_data['segmentation']
            color = (norm_feat[i] * 255).astype(np.uint8)
            mask_img[mask] = color

    # 画像合成処理
    display_frame = cv2.addWeighted(frame_rgb, ORIGINAL_WEIGHT, mask_img, MASK_WEIGHT, 0)
    result_frame = cv2.cvtColor(display_frame, cv2.COLOR_RGB2BGR)

    # 検出物体数を画面に表示
    n_masks = len(masks) if masks else 0
    text = f'検出物体数: {n_masks}'

    # Pillowで日本語テキスト描画
    font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
    img_pil = Image.fromarray(cv2.cvtColor(result_frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)

    # テキスト背景用の矩形描画
    text_bbox = draw.textbbox((10, 10), text, font=font)
    draw.rectangle([text_bbox[0]-5, text_bbox[1]-5, text_bbox[2]+5, text_bbox[3]+5],
                   fill=(0, 0, 0, 128))
    draw.text((10, 10), text, font=font, fill=(0, 255, 0))

    result_frame = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

    return result_frame, f'検出物体数: {n_masks}', current_time

print("0: 動画ファイル")
print("1: カメラ")
print("2: サンプル動画")

choice = input("選択: ")

if choice == '0':
    root = tk.Tk()
    root.withdraw()
    path = filedialog.askopenfilename()
    if not path:
        exit()
    cap = cv2.VideoCapture(path)
elif choice == '1':
    cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
    if not cap.isOpened():
        cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
else:
    # サンプル動画ダウンロード・処理
    SAMPLE_URL = 'https://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.avi'
    SAMPLE_FILE = 'vtest.avi'
    urllib.request.urlretrieve(SAMPLE_URL, SAMPLE_FILE)
    cap = cv2.VideoCapture(SAMPLE_FILE)

if not cap.isOpened():
    print('動画ファイル・カメラを開けませんでした')
    exit()

# メイン処理
print('\n=== 動画処理開始 ===')
print('操作方法:')
print('  q キー: プログラム終了')
try:
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        MAIN_FUNC_DESC = "SAMセグメンテーション結果"
        processed_frame, result, current_time = video_frame_processing(frame)
        cv2.imshow(MAIN_FUNC_DESC, processed_frame)
        if choice == '1':  # カメラの場合
            print(datetime.fromtimestamp(current_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3], result)
        else:  # 動画ファイルの場合
            print(frame_count, result)
        results_log.append(result)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
finally:
    print('\n=== プログラム終了 ===')
    cap.release()
    cv2.destroyAllWindows()
    if results_log:
        with open('result.txt', 'w', encoding='utf-8') as f:
            f.write('=== 結果 ===\n')
            f.write(f'処理フレーム数: {frame_count}\n')
            f.write(f'使用デバイス: {str(device).upper()}\n')
            if device.type == 'cuda':
                f.write(f'GPU: {torch.cuda.get_device_name(0)}\n')
            f.write('\n')
            f.write('\n'.join(results_log))
        print(f'\n処理結果をresult.txtに保存しました')