YOLOv11による2次元姿勢推定

# プログラム名: YOLOv11姿勢推定プログラム
# 特徴技術名: YOLOv11 (You Only Look Once version 11)
# 出典: Ultralytics. (2024). YOLOv11: Enhanced Feature Extraction and Optimized Efficiency. https://github.com/ultralytics/ultralytics
# 特徴機能: Ultralyticsライブラリによる姿勢推定機能を利用
# 学習済みモデル: yolo11n-pose.pt（COCOデータセット17キーポイント学習済み、2.9Mパラメータ）URL: Ultralyticsライブラリが自動ダウンロード
# 方式設計:
#   - 関連利用技術: OpenCV（画像処理・表示）、NumPy（数値計算）、Pillow（日本語フォント描画）、PyTorch（GPU/CPU自動選択）
#   - 入力と出力: 入力: 動画（ユーザは「0:動画ファイル，1:カメラ，2:サンプル動画」のメニューで選択．0:動画ファイルの場合はtkinterでファイル選択．1の場合はOpenCVでカメラが開く．2の場合はhttps://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.aviを使用）、出力: OpenCV画面でリアルタイム表示（検出された姿勢とキーポイント座標）、各フレームごとにprint()による処理結果表示、終了時にresult.txtファイルへ保存
#   - 処理手順: 1.YOLOv11モデルロード、2.動画フレーム取得、3.Ultralyticsライブラリによる姿勢推定、4.17キーポイント検出、5.信頼度閾値によるフィルタリング、6.結果の可視化
#   - 前処理、後処理: 前処理：適切なバックエンドによる低遅延フレーム取得、バッファサイズ1設定。後処理：信頼度0.5以上のキーポイントのみ表示、キーポイント名と座標の日本語フォント描画
#   - 調整を必要とする設定値: CONF_THRESHOLD（キーポイント検出の信頼度閾値、デフォルト0.5、低くすると検出数増加but誤検出も増加）
# 将来方策: 検出されたキーポイント数の時系列変化を監視し、安定して検出されるキーポイントの割合からCONF_THRESHOLDを動的に調整する機能の実装
# その他の重要事項: クロスプラットフォーム対応、COCO形式17キーポイント（鼻、目、耳、肩、肘、手首、腰、膝、足首）
# 前準備: pip install ultralytics opencv-python pillow torch

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

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

# 定数定義
MODEL_NAME = 'yolo11n-pose.pt'
CONF_THRESHOLD = 0.5  # キーポイント検出の信頼度閾値（0.0-1.0）

# 人体17箇所キーポイント名称（COCO形式）
KEYPOINT_NAMES = [
    'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
    'left_wrist', 'right_wrist', 'left_hip', 'right_hip',
    'left_knee', 'right_knee', 'left_ankle', 'right_ankle'
]

# 日本語キーポイント名称
KEYPOINT_NAMES_JP = [
    '鼻', '左目', '右目', '左耳', '右耳',
    '左肩', '右肩', '左肘', '右肘',
    '左手首', '右手首', '左腰', '右腰',
    '左膝', '右膝', '左足首', '右足首'
]

# キーポイント数定義と整合性チェック
NUM_KEYPOINTS = len(KEYPOINT_NAMES)
assert NUM_KEYPOINTS == len(KEYPOINT_NAMES_JP), "KEYPOINT_NAMESとKEYPOINT_NAMES_JPの長さが不一致である"

# プログラム開始時の説明
print('=== YOLOv11姿勢推定プログラム ===')
print('概要: YOLOv11を使用してリアルタイムで人体の17箇所のキーポイントを検出します')
print('検出部位: 鼻、目(左右)、耳(左右)、肩(左右)、肘(左右)、手首(左右)、腰(左右)、膝(左右)、足首(左右)')
print('')

# YOLOv11モデルのロード
print('YOLOv11モデルをロード中...')
model = YOLO(MODEL_NAME)
model.to(device)
print(f"モデル '{MODEL_NAME}' のロードが完了しました")
print('')

# フォント設定
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE = 16
try:
    font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
except:
    font = ImageFont.load_default()

# 結果記録用
frame_count = 0
results_log = []

def draw_japanese_text(frame, text, position, font, color=(0, 255, 0)):
    """OpenCVの画像に日本語テキストを描画する関数"""
    img_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    draw.text(position, text, font=font, fill=color)
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

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

    # YOLOv11による姿勢推定実行
    result = model(frame)[0]
    processed_frame = result.plot()

    # キーポイント処理（複数人物対応）
    keypoint_count = 0
    person_count = 0

    keypoints = result.keypoints
    if keypoints is not None:
        keypoints_data = keypoints.data.detach().cpu().numpy()

        if keypoints_data.ndim == 3 and keypoints_data.shape[1] == NUM_KEYPOINTS:
            # 人数を一括取得
            person_count = int(keypoints_data.shape[0])

            # 1フレームにつき1回だけBGR↔RGB変換してPILで描画
            img_pil = Image.fromarray(cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil)

            for person_idx in range(person_count):
                for i in range(NUM_KEYPOINTS):
                    keypoint = keypoints_data[person_idx, i]
                    if len(keypoint) >= 3:
                        x, y, conf = keypoint[:3]
                        if conf > CONF_THRESHOLD:
                            text = f'P{person_idx}:{KEYPOINT_NAMES_JP[i]}({x:.0f},{y:.0f})'
                            draw.text((int(x), int(y)), text, font=font, fill=(255, 255, 255))
                            keypoint_count += 1

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

    # 結果の整形
    if person_count > 0:
        result_str = f'検出人数:{person_count}, キーポイント数:{keypoint_count}'
    else:
        result_str = '人物未検出'

    return processed_frame, result_str, 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 キー: プログラム終了')

MAIN_FUNC_DESC = "YOLOv11姿勢推定"

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

        processed_frame, result_str, 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_str)
        else:  # 動画ファイルの場合
            print(frame_count, result_str)
        results_log.append(result_str)
        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に保存しました')