YOLO11による物体検出

# プログラム名: YOLO11物体検出プログラム
# 特徴技術名: YOLO11
# 出典: Ultralytics YOLO11 (https://github.com/ultralytics/ultralytics)
# 特徴機能: YOLO11による物体検出。COCO 80クラス全体を検出
# 学習済みモデル: YOLO11 COCO事前学習済みモデル（80クラス全体）
#   - yolo11n.pt（Nanoモデル、2.6M params、mAP50-95: 39.5%）
#   - yolo11s.pt（Smallモデル、9.4M params、mAP50-95: 47.0%）
#   - yolo11m.pt（Mediumモデル、20.1M params、mAP50-95: 51.5%）
#   - yolo11l.pt（Largeモデル、25.3M params、mAP50-95: 53.4%）
#   - yolo11x.pt（Extra Largeモデル、56.9M params、mAP50-95: 54.7%）
#   - 全モデルMS COCO 2017データセット（80クラス物体検出）対応
# 方式設計:
#   - 関連利用技術:
#     - PyTorch: 深層学習フレームワーク、CUDA対応によるGPU加速
#     - OpenCV: 画像処理、カメラ制御、描画処理、動画入出力管理
#   - 入力と出力: 入力: 動画（ユーザは「0:動画ファイル，1:カメラ，2:サンプル動画」のメニューで選択．0:動画ファイルの場合はtkinterでファイル選択．1の場合はOpenCVでカメラが開く．2の場合はhttps://github.com/opencv/opencv/raw/master/samples/data/vtest.aviを使用）、出力: OpenCV画面でリアルタイム表示（検出したオブジェクトをバウンディングボックスで表示）、各フレームごとにprint()による処理結果出力、プログラム終了時にresult.txtファイルに保存
#   - 処理手順: 1.フレーム取得、2.YOLO11推論実行、3.COCO 80クラス全体の検出、4.信頼度閾値による選別、5.バウンディングボックス描画
#   - 前処理、後処理: 前処理：YOLO11内部で自動実行（640x640リサイズ、正規化）。後処理：YOLO11内部処理により重複除去を含む後処理が自動実行。信頼度による閾値フィルタリングのみ実施
#   - 追加処理: CUDA/CPU自動検出機能により、GPU搭載環境では自動的に処理速度向上。検出結果の信頼度降順ソートにより重要な検出を優先表示
#   - 調整を必要とする設定値: CONF_THRESHOLD（オブジェクト検出信頼度閾値、デフォルト0.25）- 値を上げると誤検出が減少するが検出漏れが増加
# 将来方策: CONF_THRESHOLDの動的調整機能。フレーム毎の検出数を監視し、検出数が閾値を超えた場合は信頼度を上げ、検出数が少ない場合は下げる適応的制御の実装
# その他の重要事項: Windows環境専用設計、CUDA対応GPU推奨（自動検出・CPUフォールバック機能付き）、初回実行時は学習済みモデルの自動ダウンロード
# 前準備:
#   - pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
#   - pip install ultralytics opencv-python numpy pillow

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

warnings.filterwarnings('ignore')

# ===== 設定・定数管理 =====
# モデル情報
MODEL_INFO = {
    'n': {'name': 'nano', 'desc': '小', 'params': '2.6M', 'mAP': '39.5%'},
    's': {'name': 'small', 'desc': '小', 'params': '9.4M', 'mAP': '47.0%'},
    'm': {'name': 'medium', 'desc': '中', 'params': '20.1M', 'mAP': '51.5%'},
    'l': {'name': 'large', 'desc': '大', 'params': '25.3M', 'mAP': '53.4%'},
    'x': {'name': 'extra large', 'desc': '特大', 'params': '56.9M', 'mAP': '54.7%'}
}

# COCO 80クラス名
COCO_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
    'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
    'hair drier', 'toothbrush'
]

# クラスごとの色生成（HSVからBGRに変換）
def generate_colors(num_classes):
    colors = []
    for i in range(num_classes):
        hue = int(180.0 * i / 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

CLASS_COLORS = generate_colors(len(COCO_CLASSES))

SAMPLE_VIDEO_URL = 'https://github.com/opencv/opencv/raw/master/samples/data/vtest.avi'
SAMPLE_VIDEO_FILE = 'vtest.avi'
RESULT_FILE = 'result.txt'

# カメラ設定（必要に応じて調整）
WINDOW_WIDTH = 1280
WINDOW_HEIGHT = 720
FPS = 30

# 検出パラメータ（調整可能）
CONF_THRESHOLD = 0.25  # オブジェクト検出信頼度閾値（0.0-1.0）

# フォント設定（Pillow使用）
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE_MAIN = 20
FONT_SIZE_LABEL = 18

# ユーザガイダンス
print('=== YOLO11オブジェクト検出プログラム ===')
print('概要: リアルタイムでオブジェクトを検出し、バウンディングボックスで表示します')
print('機能: YOLO11による物体検出（COCOデータセット80クラス）')
print('操作: qキーで終了')
print('出力: 各フレームごとに処理結果を表示し、終了時にresult.txtへ保存')
print()

# モデル選択
print('使用するYOLO11モデルを選択してください:')
print('n: Nano (2.6M params, mAP 39.5%)')
print('s: Small (9.4M params, mAP 47.0%)')
print('m: Medium (20.1M params, mAP 51.5%)')
print('l: Large (25.3M params, mAP 53.4%)')
print('x: Extra Large (56.9M params, mAP 54.7%)')
print()

model_choice = input('モデル選択 (n/s/m/l/x) [デフォルト: n]: ').strip().lower()
if model_choice not in ['n', 's', 'm', 'l', 'x']:
    model_choice = 'n'
    print('デフォルトのnanoモデルを使用します')

MODEL_SIZE = model_choice
MODEL_NAME = f'yolo11{MODEL_SIZE}.pt'

# システム初期化
print()
print('システム初期化中...')
start_time = time.time()

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

# YOLO11モデル初期化
try:
    print(f'YOLO11{MODEL_SIZE}モデルを初期化中...')
    model = YOLO(MODEL_NAME)
    model.to(device)
    print(f'YOLO11{MODEL_SIZE}モデルの初期化が完了しました')
    print(f'モデル: {MODEL_INFO[MODEL_SIZE]["name"]} ({MODEL_INFO[MODEL_SIZE]["params"]} params, mAP {MODEL_INFO[MODEL_SIZE]["mAP"]})')
except Exception as e:
    print(f'YOLO11{MODEL_SIZE}モデルの初期化に失敗しました')
    print(f'エラー: {e}')
    exit()

print(f'{str(device).upper()}使用モード')
print('初期化完了')
print()

# グローバル変数
frame_count = 0
results_log = []

def draw_text_pil(bgr_frame, texts):
    """
    Pillow + Meiryoで日本語テキストを描画する。
    texts: list of dict with keys {text, org, color}
    """
    font_main = ImageFont.truetype(FONT_PATH, FONT_SIZE_MAIN)
    font_label = ImageFont.truetype(FONT_PATH, FONT_SIZE_LABEL)
    img_pil = Image.fromarray(cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    for item in texts:
        text = item['text']
        x, y = item['org']
        color = item['color']  # RGB
        font = font_main if item.get('font', 'main') == 'main' else font_label
        draw.text((x, y), text, font=font, fill=color)
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

def video_frame_processing(frame):
    """
    フレーム処理
    - 推論実行
    - バウンディングボックス描画
    - 結果文字列の生成（B案: 各検出ボックスのクラス、信頼度、座標をCSV風で列挙）
    """
    global frame_count
    current_time = time.time()
    frame_count += 1

    # 推論実行（閾値はmodel側で適用。後段の二重判定は行わない）
    results = model(frame, conf=CONF_THRESHOLD, verbose=False)
    objects = []

    r0 = results[0]
    if r0.boxes is not None and len(r0.boxes) > 0:
        boxes = r0.boxes.xyxy.cpu().numpy()
        confs = r0.boxes.conf.cpu().numpy()
        classes = r0.boxes.cls.cpu().numpy()

        # 信頼度でソート（降順）
        sorted_indices = np.argsort(confs)[::-1]
        boxes = boxes[sorted_indices]
        confs = confs[sorted_indices]
        classes = classes[sorted_indices]

        for box, conf, cls in zip(boxes, confs, classes):
            x1, y1, x2, y2 = map(int, box)
            class_id = int(cls)
            class_name = COCO_CLASSES[class_id] if 0 <= class_id < len(COCO_CLASSES) else str(class_id)
            objects.append({
                'box': (x1, y1, x2, y2),
                'conf': float(conf),
                'class_id': class_id,
                'class_name': class_name
            })

    # 描画処理（矩形はOpenCV、テキストはPillow）
    for obj in objects:
        x1, y1, x2, y2 = obj['box']
        class_id = obj['class_id']
        color = CLASS_COLORS[class_id] if 0 <= class_id < len(CLASS_COLORS) else (0, 255, 0)
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

    # ラベル・情報テキスト（Pillowで日本語表示）
    text_items = []
    for obj in objects:
        x1, y1, x2, y2 = obj['box']
        class_name = obj['class_name']
        conf_pct = f'{obj["conf"]*100:.1f}%'
        # RGB色に変換
        b, g, r = CLASS_COLORS[obj['class_id']] if 0 <= obj['class_id'] < len(CLASS_COLORS) else (0, 255, 0)
        text_items.append({'text': f'{class_name}', 'org': (x1, max(0, y1 - 22)), 'color': (r, g, b), 'font': 'label'})
        text_items.append({'text': f'Conf:{conf_pct}', 'org': (x1, min(y2 + 5, frame.shape[0]-20)), 'color': (255, 255, 255), 'font': 'label'})

    info1 = f'YOLO11{MODEL_SIZE} ({str(device).upper()}) | Frame: {frame_count} | Objects: {len(objects)}'
    info2 = '操作: q キーで終了'
    text_items.append({'text': info1, 'org': (10, 10), 'color': (255, 255, 255), 'font': 'main'})
    text_items.append({'text': info2, 'org': (10, 35), 'color': (255, 255, 0), 'font': 'label'})

    frame = draw_text_pil(frame, text_items)

    # 出力用の結果文字列（B案: CSV風列挙）
    if len(objects) == 0:
        result = 'count=0'
    else:
        parts = []
        for obj in objects:
            x1, y1, x2, y2 = obj['box']
            parts.append(f'class={obj["class_name"]},conf={obj["conf"]:.3f},box=[{x1},{y1},{x2},{y2}]')
        result = f'count={len(objects)}; ' + ' | '.join(parts)

    return frame, result, current_time

# 入力選択
print('入力ソースを選択してください:')
print('0: 動画ファイル')
print('1: カメラ')
print('2: サンプル動画')
print()

choice = input('選択 (0/1/2): ').strip()
temp_file = None

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)
    # 必要に応じて解像度設定（任意）
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WINDOW_WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, WINDOW_HEIGHT)
    cap.set(cv2.CAP_PROP_FPS, FPS)
else:
    # サンプル動画ダウンロード・処理
    try:
        urllib.request.urlretrieve(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_FILE)
        temp_file = SAMPLE_VIDEO_FILE
    except Exception as e:
        print(f"動画のダウンロードに失敗しました: {SAMPLE_VIDEO_URL}")
        print(f"エラー: {e}")
        exit()
    cap = cv2.VideoCapture(SAMPLE_VIDEO_FILE)

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

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

MAIN_FUNC_DESC = "YOLO11 物体検出"
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('=== YOLO11オブジェクト検出結果 ===\n')
            f.write(f'処理フレーム数: {frame_count}\n')
            f.write(f'使用モデル: YOLO11{MODEL_SIZE} ({MODEL_INFO[MODEL_SIZE]["name"]})\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}に保存しました')

    if temp_file and os.path.exists(temp_file):
        os.remove(temp_file)