YOLOv8による2次元姿勢推定

# プログラム名: YOLOv8リアルタイム姿勢推定
# 特徴技術名: YOLOv8（You Only Look Once version 8）
# 出典: Jocher, G., Chaurasia, A., & Qiu, J. (2023). YOLO by Ultralytics. GitHub. https://github.com/ultralytics/ultralytics
# 特徴機能: リアルタイム姿勢推定 - 単一のニューラルネットワークで人体の17個のキーポイントを検出
# 学習済みモデル: YOLOv8姿勢推定モデル（yolov8n-pose.pt）- COCOデータセットで事前学習済み
# 方式設計:
#   - 関連利用技術: OpenCV（動画処理・表示）、NumPy（数値計算）、tkinter（ファイル選択ダイアログ）、Pillow（日本語表示）
#   - 入力と出力: 動画入力（ファイル/カメラ/サンプル）、OpenCV画面表示、result.txt保存
#   - 処理手順: 1.動画フレーム取得、2.YOLOv8推論、3.キーポイント抽出、4.可視化、5.結果表示
#   - 前処理、後処理: フレーム取得、信頼度フィルタリング
#   - 追加処理: キーポイント名称の日本語表示
#   - 調整を必要とする設定値: CONF_THRESHOLD（人物検出閾値）、KPT_THRESHOLD（キーポイント信頼度閾値）

import cv2
import numpy as np
from ultralytics import YOLO
import tkinter as tk
from tkinter import filedialog
import time
import urllib.request
import torch
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime

# YOLOv8設定
MODEL_NAME = 'yolov8n-pose.pt'
CONF_THRESHOLD = 0.5   # 人物検出の信頼度閾値
KPT_THRESHOLD = 0.25   # キーポイント検出の信頼度閾値

# 表示設定
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE = 14
TEXT_COLOR_RGB = (0, 255, 0)  # PIL用RGB形式

KEYPOINT_NAMES = [
    '鼻', '左目', '右目', '左耳', '右耳',
    '左肩', '右肩', '左肘', '右肘',
    '左手首', '右手首', '左腰', '右腰',
    '左膝', '右膝', '左足首', '右足首'
]

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

# プログラム開始メッセージ
print('='*60)
print('YOLOv8リアルタイム姿勢推定プログラム')
print('='*60)
print('概要説明:')
print('  動画から人体の17個のキーポイントを検出し、リアルタイム表示します')
print()
print('操作方法:')
print('  - メニューから入力ソースを選択（0:ファイル, 1:カメラ, 2:サンプル）')
print('  - qキー: プログラムを終了')
print()
print('注意事項:')
print('  - 検出結果は各フレームごとにコンソール表示')
print('  - 終了時にresult.txtへ自動保存')
print('='*60)
print()

# フォント初期化
try:
    font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
except:
    font = None
    print('警告: 日本語フォントの読み込みに失敗しました')

# モデルロード
print('YOLOv8モデルをロード中...')
try:
    model = YOLO(MODEL_NAME)
    model.to(device)
    print(f'モデル {MODEL_NAME} をロードしました')
except Exception as e:
    print(f'モデルのロードに失敗しました: {e}')
    exit()

def draw_japanese_text(img, text, position, color_rgb):
    """OpenCV画像に日本語テキストを描画"""
    if font is None:
        return img

    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    draw.text(position, text, font=font, fill=color_rgb)
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

frame_count = 0
results_log = []

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

    # 推論実行
    results = model(frame, verbose=False, conf=CONF_THRESHOLD)

    # 注釈付きフレームを取得
    annotated_frame = results[0].plot()

    detected_keypoints = []
    result_text = ""

    for result in results:
        if hasattr(result, 'keypoints') and result.keypoints is not None:
            keypoints = result.keypoints

            if hasattr(keypoints, 'xy'):
                # xy属性から座標を取得（画像座標系）
                keypoints_xy = keypoints.xy.cpu().numpy() if hasattr(keypoints.xy, 'cpu') else keypoints.xy
                keypoints_conf = keypoints.conf.cpu().numpy() if hasattr(keypoints.conf, 'cpu') else keypoints.conf

                # 各人物のキーポイントを処理
                for person_idx in range(keypoints_xy.shape[0]):
                    person_keypoints = []
                    for i in range(17):
                        if keypoints_conf[person_idx, i] > KPT_THRESHOLD:
                            x, y = keypoints_xy[person_idx, i]
                            conf = keypoints_conf[person_idx, i]

                            # 座標が有効な範囲内にあるか確認
                            if x > 0 and y > 0:
                                text = f'{KEYPOINT_NAMES[i]}'
                                annotated_frame = draw_japanese_text(
                                    annotated_frame, text,
                                    (int(x), int(y)), TEXT_COLOR_RGB
                                )
                                person_keypoints.append(
                                    f'{KEYPOINT_NAMES[i]}: ({x:.0f}, {y:.0f}) [信頼度: {conf:.2f}]'
                                )

                    if person_keypoints:
                        detected_keypoints.append(person_keypoints)

    # 結果テキストの生成
    if detected_keypoints:
        result_parts = []
        for person_idx, person_keypoints in enumerate(detected_keypoints):
            result_parts.append(f'人物{person_idx + 1}: {len(person_keypoints)}個のキーポイント検出')
        result_text = ', '.join(result_parts)
    else:
        result_text = '検出なし'

    return annotated_frame, result_text, current_time

print('入力ソースを選択してください:')
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 = "YOLOv8姿勢推定"
        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に保存しました')