YOLO11によるOBB（回転物体検出）

# プログラム名: YOLO11-OBB物体検出プログラム
# 特徴技術名: YOLO11-OBB
# 出典:
#   - 論文: Khanam, R., & Hussain, M. (2024). YOLOv11: An Overview of the Key Architectural Enhancements. arXiv preprint arXiv:2410.17725.
#   - DOTA論文: Xia, G. S., et al. (2018). DOTA: A Large-scale Dataset for Object Detection in Aerial Images. CVPR 2018.
#   - GitHub: https://github.com/ultralytics/ultralytics
#   - 公式ドキュメント: https://docs.ultralytics.com/tasks/obb/
# 特徴機能: Oriented Bounding Box (OBB) による任意角度物体検出機能。従来の軸平行境界ボックスと異なり、物体の向きに合わせて回転可能な境界ボックスによる検出。航空・衛星画像において物体が様々な角度で現れる場合に有効で、背景領域の誤検出を削減
# 学習済みモデル:
#   - yolo11n-obb.pt（DOTAv1データセット15クラス対応、Nano版、最軽量、自動ダウンロード）
#   - yolo11s-obb.pt（DOTAv1データセット15クラス対応、Small版、軽量、自動ダウンロード）
#   - yolo11m-obb.pt（DOTAv1データセット15クラス対応、Medium版、バランス型、自動ダウンロード）
#   - yolo11l-obb.pt（DOTAv1データセット15クラス対応、Large版、精度重視、自動ダウンロード）
#   - yolo11x-obb.pt（DOTAv1データセット15クラス対応、Extra Large版、最高精度、自動ダウンロード）
# 方式設計:
#   - 関連利用技術:
#     - PyTorch: 深層学習フレームワーク、CUDA対応によるGPU加速
#     - OpenCV: 画像処理、カメラ制御、描画処理、動画入出力管理
#     - Matplotlib: 学習曲線の可視化
#   - 入力と出力: 入力: 動画（ユーザは「0:動画ファイル，1:カメラ，2:サンプル動画」のメニューで選択。0:動画ファイルの場合はtkinterでファイル選択。1の場合はOpenCVでカメラが開く。2の場合はhttps://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.aviを使用）、出力: OpenCV画面でリアルタイム表示（検出したオブジェクトをOBBで表示）、各フレームごとにprint()による処理結果表示、プログラム終了時にresult.txtファイルに保存、学習ログが存在する場合は自動的に学習曲線を画像ファイルとして保存
#   - 処理手順: 1.学習ログの自動検索と可視化、2.フレーム取得、3.YOLO11-OBB推論実行、4.DOTAv1 15クラスの検出、5.信頼度閾値による選別、6.Oriented Bounding Box描画
#   - 前処理、後処理: 前処理：YOLO11-OBB内部で自動実行（640x640リサイズ、正規化）。後処理：YOLO11-OBBのNMSフリー設計により、重複除去処理が不要。信頼度による閾値フィルタリングのみ実施
#   - 追加処理: CUDA/CPU自動検出機能により、GPU搭載環境では自動的に加速。検出結果の信頼度降順ソートにより重要な検出を優先表示。学習ログの自動検出と可視化
#   - 調整を必要とする設定値: CONF_THRESH（オブジェクト検出信頼度閾値、デフォルト0.5）- 値を上げると誤検出が減少するが検出漏れが増加
# 将来方策: CONF_THRESHの動的調整機能。フレーム毎の検出数を監視し、検出数が閾値を超えた場合は信頼度を上げ、検出数が少ない場合は下げる適応的制御の実装
# その他の重要事項: 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 matplotlib pandas

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 PIL import Image, ImageDraw, ImageFont
import matplotlib
matplotlib.use('Agg')  # GUIバックエンドを使用しない
import matplotlib.pyplot as plt
import pandas as pd
from pathlib import Path
import glob
from datetime import datetime

warnings.filterwarnings('ignore')

# 日本語フォント設定
FONT_PATH = 'C:/Windows/Fonts/meiryo.ttc'
FONT_SIZE = 30
FONT = ImageFont.truetype(FONT_PATH, FONT_SIZE)

# matplotlib日本語フォント設定
plt.rcParams['font.sans-serif'] = ['MS Gothic']
plt.rcParams['axes.unicode_minus'] = False

# ===== 設定・定数管理 =====
# YOLO11-OBBモデル設定（デフォルト：n、変更可能：n, s, m, l, x）
MODEL_SIZE = 'n'  # 使用するモデルサイズ（n=nano, s=small, m=medium, l=large, x=extra large）
MODEL_NAME = f'yolo11{MODEL_SIZE}-obb.pt'

# モデル情報
MODEL_INFO = {
    'n': {'name': 'nano', 'desc': '最軽量'},
    's': {'name': 'small', 'desc': '軽量'},
    'm': {'name': 'medium', 'desc': '中程度'},
    'l': {'name': 'large', 'desc': '精度重視'},
    'x': {'name': 'extra large', 'desc': '最高精度'}
}

# DOTAv1 15クラス名
DOTA_CLASSES = [
    'plane', 'ship', 'storage-tank', 'baseball-diamond', 'tennis-court', 'basketball-court',
    'ground-track-field', 'harbor', 'bridge', 'large-vehicle', 'small-vehicle',
    'helicopter', 'roundabout', 'soccer-ball-field', 'swimming-pool'
]

# クラスごとの色生成（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(DOTA_CLASSES))

# 検出パラメータ（調整可能）
CONF_THRESH = 0.5  # オブジェクト検出信頼度閾値（0.0-1.0）
IMG_SIZE = 640  # YOLO入力画像サイズ

# 表示設定
OBJECT_TEXT_Y = 90  # 検出数表示Y座標
OBJECT_COLOR = (0, 255, 0)  # 検出数表示色（BGR）


def visualize_training_curves():
    """学習ログの自動検索と可視化"""
    # 学習ログファイルの検索パターン（精度向上のため再帰検索を使用）
    search_patterns = [
        '**/runs/obb/train*/results.csv',
        '**/runs/detect/train*/results.csv',
        '**/runs/train*/results.csv',
        'runs/**/results.csv',
        '**/results.csv'
    ]

    results_file = None
    all_files = []

    # 各パターンで検索
    for pattern in search_patterns:
        files = glob.glob(pattern, recursive=True)
        all_files.extend(files)

    # 重複を除去して最新のファイルを選択
    if all_files:
        unique_files = list(set(all_files))
        results_file = max(unique_files, key=os.path.getmtime)

    if not results_file:
        print("学習ログファイルが見つかりませんでした")
        return False

    print(f"学習ログを検出: {results_file}")

    try:
        # CSVファイル読み込み
        df = pd.read_csv(results_file)
        df.columns = df.columns.str.strip()

        # 利用可能なカラムを確認（OBB特有のメトリクスも含む）
        available_metrics = []
        metric_mapping = {
            'train/box_loss': 'Box Loss (Train)',
            'train/cls_loss': 'Class Loss (Train)',
            'train/obb_loss': 'OBB Loss (Train)',
            'val/box_loss': 'Box Loss (Val)',
            'val/cls_loss': 'Class Loss (Val)',
            'val/obb_loss': 'OBB Loss (Val)',
            'metrics/mAP50': 'mAP@0.5',
            'metrics/mAP50-95': 'mAP@0.5:0.95'
        }

        for col in df.columns:
            if col in metric_mapping:
                available_metrics.append(col)

        if not available_metrics:
            print("可視化可能なメトリクスが見つかりませんでした")
            return False

        # グラフ作成
        num_plots = min(len(available_metrics), 4)
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        axes = axes.flatten()

        for i, metric in enumerate(available_metrics[:4]):
            axes[i].plot(df.index + 1, df[metric], linewidth=2)
            axes[i].set_title(metric_mapping.get(metric, metric))
            axes[i].set_xlabel('エポック')
            axes[i].set_ylabel('値')
            axes[i].grid(True, alpha=0.3)

        # 未使用のサブプロットを非表示
        for i in range(num_plots, 4):
            axes[i].set_visible(False)

        plt.suptitle('YOLO11-OBB 学習曲線', fontsize=16)
        plt.tight_layout()

        # グラフを保存（表示はしない）
        output_path = 'training_curves.png'
        plt.savefig(output_path, dpi=100)
        plt.close()  # メモリ解放
        print(f"学習曲線を保存: {output_path}")

        # 最終エポックの統計情報
        print("\n=== 最終エポックの統計 ===")
        last_epoch = df.iloc[-1]

        for metric in available_metrics:
            value = last_epoch[metric]
            print(f"{metric_mapping.get(metric, metric)}: {value:.4f}")

        return True

    except Exception as e:
        print(f"学習ログの可視化中にエラー: {e}")
        return False


def add_text(frame, text, y, color):
    """フレームに日本語テキストを追加する関数"""
    img_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    draw.text((10, y), text, font=FONT, fill=(color[2], color[1], color[0]))
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)


# プログラム概要表示
print('=== YOLO11-OBBオブジェクト検出プログラム ===')
print('概要: リアルタイムでオブジェクトを検出し、Oriented Bounding Boxで表示します')
print('機能: YOLO11-OBBによるオブジェクト検出（DOTAv1データセット15クラス）')
print('操作: qキーで終了')
print('出力: 各フレームごとの処理結果表示、終了時にresult.txt保存')
print()

# 学習ログの自動可視化
print('学習ログを検索中...')
visualize_training_curves()
print()

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

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

if device.type == 'cuda':
    print(f'GPU検出: {torch.cuda.get_device_name(0)}')
    print(f'CUDA バージョン: {torch.version.cuda}')

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

print('初期化完了')
print()

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


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

    # オブジェクト検出実行（デバイス指定）
    preds = model(frame, conf=CONF_THRESH, device=device,
                  imgsz=IMG_SIZE, verbose=False)

    # 標準の描画を使用
    frame_out = preds[0].plot()

    # 検出結果処理（ソート用のコピーを作成）
    objects = []
    for pred in preds:
        if hasattr(pred, 'obb') and pred.obb is not None and len(pred.obb.cls) > 0:
            obj_cnt = len(pred.obb.cls)
            frame_out = add_text(frame_out, f'検出数: {obj_cnt}', OBJECT_TEXT_Y, OBJECT_COLOR)

            # データをコピーしてからソート
            xyxyxyxy = pred.obb.xyxyxyxy.cpu().numpy().copy()
            confs = pred.obb.conf.cpu().numpy().copy()
            classes = pred.obb.cls.cpu().numpy().copy()

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

            # 各オブジェクトの処理
            for i, (obb, conf, cls) in enumerate(zip(xyxyxyxy, confs, classes)):
                class_id = int(cls)
                object_data = {
                    'obb': obb,
                    'detection_conf': conf,
                    'class_id': class_id,
                    'class_name': DOTA_CLASSES[class_id]
                }
                objects.append(object_data)

    # 結果文字列の作成
    if objects:
        # クラス別検出数をカウント
        class_counts = {}
        for obj in objects:
            class_name = obj['class_name']
            class_counts[class_name] = class_counts.get(class_name, 0) + 1

        result = f'{len(objects)}個検出'
        # 信頼度順に上位3件の詳細を表示
        for i, obj in enumerate(objects[:3]):
            result += f' | Top{i+1}: {obj["class_name"]}({obj["detection_conf"]:.2f})'
    else:
        result = '検出なし'

    # システム情報表示
    info1 = f'YOLO11-OBB ({str(device)}) | フレーム: {frame_count} | オブジェクト: {len(objects)}'
    info2 = 'キー操作: q=終了'

    frame_out = add_text(frame_out, info1, 10, (255, 255, 255))
    frame_out = add_text(frame_out, info2, 50, (255, 255, 0))

    return frame_out, result, 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 = "YOLO11-OBB Object Detection"
        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に保存しました')