RT-DETRv2による物体検出

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

# RT-DETRv2物体検出プログラム
# 特徴技術名: RT-DETRv2 (Real-Time Detection Transformer version 2)
# 出典: W. Lv, Y. Zhao, Q. Chang, K. Huang, G. Wang, and Y. Liu, "RT-DETRv2: Improved Baseline with Bag-of-Freebies for Real-Time Detection Transformer," arXiv preprint arXiv:2407.17140, 2024.
# 特徴機能: Bag-of-Freebiesによる改良ベースラインと選択的マルチスケール特徴抽出による物体検出
# 学習済みモデル: PekingU/rtdetr_v2_r50vd（Hugging Face Transformers）
# 方式設計:
#   - 関連利用技術:
#     - PyTorch: 深層学習フレームワーク、CUDA対応
#     - Transformers: Hugging Face Transformersライブラリ
#     - OpenCV: 画像処理、カメラ制御、描画処理、動画入出力管理
#   - 入力と出力:
#     入力: 動画（0:動画ファイル，1:カメラ，2:サンプル動画）
#     出力: OpenCV画面で表示（バウンディングボックス表示）、各フレームごとにprint()で処理結果を表示、終了時にresult.txtへ保存
#   - 処理手順: 1.フレーム取得、2.RT-DETRv2推論、3.COCO系クラス検出、4.信頼度閾値で選別、5.バウンディングボックス描画
#   - 前処理・後処理: 前処理はプロセッサで画像をテンソル化、後処理はpost_process_object_detection()による結果整形
#   - 追加処理: GPU/CPU自動選択、検出結果のスコア降順ソート
#   - 調整可能値: CONF_THRESHOLD（検出信頼度閾値、デフォルト0.25）
# 重要事項: Windows環境での動作を想定（フォントはC:/Windows/Fonts/meiryo.ttc を使用）

import os
import cv2
import time
import urllib.request
import warnings
import tkinter as tk
from tkinter import filedialog
from datetime import datetime

import torch
import numpy as np
from transformers import RTDetrV2ForObjectDetection, AutoImageProcessor
from PIL import Image, ImageFont, ImageDraw

warnings.filterwarnings('ignore')

# ===== 設定・定数 =====
MODEL_NAME = 'PekingU/rtdetr_v2_r50vd'
CONF_THRESHOLD = 0.25
SAMPLE_VIDEO_URL = 'https://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.avi'
SAMPLE_VIDEO_FILE = 'vtest.avi'
RESULT_FILE = 'result.txt'

# OpenCVウィンドウ内テキスト表示に使用するフォント（Windows想定）
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE = 20

# クラスごとの色生成（HSV->BGR）
def generate_colors(num_classes: int):
    colors = []
    for i in range(num_classes):
        hue = int(180.0 * i / max(1, num_classes))
        hsv = np.uint8([[[hue, 255, 255]]])
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0][0]
        colors.append((int(bgr[0]), int(bgr[1]), int(bgr[2])))
    return colors

# ===== ガイダンス表示 =====
print('=== RT-DETRv2物体検出プログラム ===')
print('概要: RT-DETRv2による物体検出を行い、検出結果を各フレームごとに表示する')
print('操作方法: 0=動画ファイル, 1=カメラ, 2=サンプル動画 / qキーで終了')
print('注意: Windows環境を想定（フォントにMeiryoを使用）')
print()

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

# ===== モデル・プロセッサ初期化 =====
try:
    print('モデルを初期化しています...')
    model = RTDetrV2ForObjectDetection.from_pretrained(MODEL_NAME)
    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
    model.to(device)
    print('モデル初期化が完了しました')
except Exception as e:
    print('モデル初期化に失敗しました')
    print(f'エラー: {e}')
    raise SystemExit(1)

# ラベルマッピング（id2label）とクラス色
# id2labelのキーはint想定
id2label = {int(k): v for k, v in model.config.id2label.items()} if isinstance(list(model.config.id2label.keys())[0], (str, int)) else model.config.id2label
NUM_CLASSES = len(id2label)
CLASS_COLORS = generate_colors(NUM_CLASSES)

# ===== グローバル状態 =====
frame_count = 0
results_log = []

# ===== テキスト描画（Pillow）補助 =====
def draw_texts_with_pillow(bgr_image, texts):
    """
    texts: list of dict { 'text': str, 'org': (x,y), 'color_bgr': (b,g,r) }
    """
    try:
        font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
    except Exception:
        # 指定フォントが読み込めない場合はデフォルトフォント
        font = ImageFont.load_default()

    img_pil = Image.fromarray(cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)

    for t in texts:
        text = t['text']
        org = t['org']
        b, g, r = t['color_bgr']
        # PillowはRGB指定
        draw.text(org, text, font=font, fill=(r, g, b))

    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

# ===== メイン処理関数（1フレーム） =====
def video_frame_processing(frame):
    global frame_count
    current_time = time.time()
    frame_count += 1

    # 前処理
    frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    inputs = processor(images=frame_pil, return_tensors='pt')
    # 入力テンソルをデバイスへ
    for k in inputs:
        if isinstance(inputs[k], torch.Tensor):
            inputs[k] = inputs[k].to(device)

    # 推論
    with torch.no_grad():
        outputs = model(**inputs)

    # 後処理（thresholdはpost_process側に一本化）
    target_sizes = torch.tensor([frame.shape[:2]], device=device)
    results = processor.post_process_object_detection(
        outputs,
        target_sizes=target_sizes,
        threshold=CONF_THRESHOLD
    )[0]

    # 結果の整形と描画
    result_items = []
    texts_to_draw = []

    if len(results['boxes']) > 0:
        boxes = results['boxes'].detach().cpu().numpy()
        scores = results['scores'].detach().cpu().numpy()
        labels = results['labels'].detach().cpu().numpy()

        # スコア降順
        order = np.argsort(scores)[::-1]
        boxes = boxes[order]
        scores = scores[order]
        labels = labels[order]

        # バウンディングボックス描画（枠はOpenCV、テキストはPillow）
        for box, score, label_id in zip(boxes, scores, labels):
            x1, y1, x2, y2 = map(int, box)
            class_name = id2label.get(int(label_id), str(int(label_id)))
            color = CLASS_COLORS[int(label_id) % NUM_CLASSES]

            # 枠
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

            # ラベルとスコア（Pillowで描画）
            texts_to_draw.append({
                'text': f'{class_name}',
                'org': (x1, max(0, y1 - 10)),
                'color_bgr': color
            })
            texts_to_draw.append({
                'text': f'Conf:{score:.1%}',
                'org': (x1, min(frame.shape[0] - 5, y2 + 15)),
                'color_bgr': (255, 255, 255)
            })

            # ログ用の整形
            result_items.append(f'{class_name}, score={score:.3f}, box=[{x1},{y1},{x2},{y2}]')

    # 画面上部情報（Pillowで描画）
    texts_to_draw.append({
        'text': f'RT-DETRv2 ({str(device).upper()})  Frame:{frame_count}  Objects:{len(result_items)}',
        'org': (10, 30),
        'color_bgr': (255, 255, 255)
    })
    texts_to_draw.append({
        'text': 'q: 終了',
        'org': (10, 60),
        'color_bgr': (255, 255, 0)
    })

    frame = draw_texts_with_pillow(frame, texts_to_draw)

    # 結果テキスト（各検出の座標・スコア列挙）
    if result_items:
        result_text = '; '.join(result_items)
    else:
        result_text = 'no detections'

    return frame, result_text, current_time

# ===== 入力選択 =====
print('0: 動画ファイル')
print('1: カメラ')
print('2: サンプル動画')

choice = input('選択: ').strip()

if choice == '0':
    root = tk.Tk()
    root.withdraw()
    path = filedialog.askopenfilename()
    if not path:
        raise SystemExit(0)
    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:
    # サンプル動画ダウンロード・処理
    try:
        urllib.request.urlretrieve(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_FILE)
    except Exception as e:
        print('サンプル動画のダウンロードに失敗しました')
        print(f'エラー: {e}')
        raise SystemExit(1)
    cap = cv2.VideoCapture(SAMPLE_VIDEO_FILE)

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

# ===== メイン処理 =====
print('\n=== 動画処理開始 ===')
print('操作方法:')
print('  q キー: プログラム終了')

MAIN_FUNC_DESC = 'RT-DETRv2 Object Detection'

try:
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        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_FILE, 'w', encoding='utf-8') as f:
            f.write('=== RT-DETRv2オブジェクト検出結果 ===\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_FILE}に保存しました')