YOLOE によるオープンボキャブラリ物体検出・追跡

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

REM Git をシステム領域にインストール
winget install --scope machine --id Git.Git -e --silent
REM Git のパス設定
set "GIT_PATH=C:\Program Files\Git\cmd"
if exist "%GIT_PATH%" (
    echo "%PATH%" | find /i "%GIT_PATH%" >nul
    if errorlevel 1 setx PATH "%PATH%;%GIT_PATH%" /M >nul
)

# プログラム名: YOLOE オープンボキャブラリ物体検出・追跡プログラム
# 特徴技術名: YOLOE (Real-Time Seeing Anything)
# 出典: Ultralytics YOLOE公式ドキュメント - https://docs.ultralytics.com/models/yoloe/
# 特徴機能: オープンボキャブラリ物体検出とセグメンテーション機能。テキスト、視覚、内部ボキャブラリプロンプトによる任意物体検出・追跡が可能
# 学習済みモデル: yoloe-11s-seg.pt (小型セグメンテーション), yoloe-11m-seg.pt (中型), yoloe-11l-seg.pt (大型) - Objects365 + LVIS データセット対応、Ultralyticsから自動ダウンロード
# 方式設計:
#   - 関連利用技術:
#     - PyTorch: 深層学習フレームワーク、CUDA対応によるGPU加速
#     - OpenCV: 画像処理、カメラ制御、描画処理、動画入出力管理
#   - 入力と出力: 入力: 動画（ユーザは「0:動画ファイル，1:カメラ，2:サンプル動画」のメニューで選択．0:動画ファイルの場合はtkinterでファイル選択．1の場合はOpenCVでカメラが開く．2の場合はhttps://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.aviを使用）、出力: OpenCV画面でリアルタイム表示（検出したオブジェクトをバウンディングボックスとセグメンテーションマスクで表示、物体ID・軌跡・滞在時間も表示）、1秒間隔でprint()による処理結果表示（物体ID・滞在時間・移動距離を含む）、プログラム終了時にresult.txtファイルに保存
#   - 処理手順: 1.フレーム取得、2.YOLOE推論実行、3.オープンボキャブラリによる任意物体検出・セグメンテーション、4.フレーム間物体追跡・ID割り当て、5.信頼度閾値による選別、6.バウンディングボックス・セグメンテーションマスク・追跡情報描画
#   - 前処理、後処理: 前処理：YOLOE内部で自動実行（640x640リサイズ、正規化）。後処理：オープンボキャブラリ機能による領域テキストアライメント処理。セグメンテーションマスクの重畳描画。物体追跡による継続性維持。信頼度による閾値フィルタリング実施
#   - 追加処理: CUDA/CPU自動検出機能により、GPU搭載環境では自動的に加速。検出結果の信頼度ソートにより重要な検出を優先表示。セグメンテーションマスクの色分け表示。フレーム間物体追跡による軌跡記録・滞在時間計測・移動距離計算
#   - 調整を必要とする設定値: CONF_THRESHOLD（オブジェクト検出信頼度閾値、デフォルト0.3）- 値を上げると誤検出が減少するが検出漏れが増加
# 将来方策: 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

import cv2
import tkinter as tk
from tkinter import filedialog
import os
import torch
import numpy as np
from ultralytics import YOLOE
import warnings
import time
import urllib.request
import math

warnings.filterwarnings('ignore')

# ===== 設定・定数管理 =====
# URLとファイル名
SAMPLE_VIDEO_URL = 'https://raw.githubusercontent.com/opencv/opencv/master/samples/data/vtest.avi'
SAMPLE_VIDEO_FILE = 'vtest.avi'

# ウィンドウ名
WINDOW_NAME = 'YOLOE Object Detection & Tracking'

# カスタムクラス定義（オープンボキャブラリの例）
CLASSES = [
    'person wearing red shirt', 'person wearing black', 'small car', 'big car',
    'small dog', 'large dog', 'orange cat', 'black cat', 'flying bird', 'perched bird',
    'open laptop', 'closed laptop', 'smartphone', 'tablet', 'thick book', 'thin book',
    'coffee cup', 'tea cup', 'water bottle', 'plastic bottle', 'wooden chair', 'metal chair',
    'dining table', 'computer desk', 'flat screen tv', 'old tv', 'wireless keyboard', 'gaming mouse',
    'red apple', 'green apple', 'fresh orange', 'ripe banana', 'hot pizza', 'birthday cake',
    'colorful flower', 'white flower', 'green tree', 'dry tree', 'school backpack', 'hiking backpack'
]

RESULT_FILE = 'result.txt'

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

# 追跡パラメータ（調整可能）
MAX_TRACKING_DISTANCE = 100  # 追跡可能最大距離（ピクセル）
MAX_TRACKING_DISTANCE_SQ = MAX_TRACKING_DISTANCE ** 2  # 距離の2乗（高速化用）
MAX_LOST_FRAMES = 30  # 追跡失敗許容フレーム数
TRAIL_LENGTH = 10  # 軌跡表示フレーム数

# 表示設定
PRINT_INTERVAL = 1.0  # 結果出力間隔（秒）


class ObjectTracker:
    def __init__(self):
        self.next_id = 1
        self.tracked_objects = {}
        self.frame_count = 0

    def calculate_distance_squared(self, pos1, pos2):
        # 平方根計算を省略して高速化
        return (pos1[0] - pos2[0])**2 + (pos1[1] - pos2[1])**2

    def update(self, detections):
        self.frame_count += 1
        current_positions = []
        current_classes = []
        current_scores = []
        current_boxes = []
        current_masks = []

        # 検出結果から位置とクラス情報を抽出
        for detection in detections:
            if detection.boxes is not None and len(detection.boxes) > 0:
                boxes = detection.boxes.xyxy.cpu().numpy()
                scores = detection.boxes.conf.cpu().numpy()
                classes = detection.boxes.cls.cpu().numpy()
                masks = detection.masks.data.cpu().numpy() if detection.masks is not None else None

                for i, (box, score, cls) in enumerate(zip(boxes, scores, classes)):
                    center_x = (box[0] + box[2]) / 2
                    center_y = (box[1] + box[3]) / 2
                    current_positions.append((center_x, center_y))
                    current_classes.append(int(cls))
                    current_scores.append(score)
                    current_boxes.append(box)
                    if masks is not None and i < len(masks):
                        current_masks.append(masks[i])
                    else:
                        current_masks.append(None)

        # 既存の追跡オブジェクトとのマッチング
        matched_objects = {}
        used_detections = set()

        for obj_id, obj_data in self.tracked_objects.items():
            best_match = -1
            best_distance_sq = float('inf')

            for i, pos in enumerate(current_positions):
                if i in used_detections:
                    continue
                distance_sq = self.calculate_distance_squared(obj_data['last_position'], pos)
                if distance_sq < MAX_TRACKING_DISTANCE_SQ and distance_sq < best_distance_sq:
                    best_distance_sq = distance_sq
                    best_match = i

            if best_match != -1:
                # マッチングした場合、情報を更新
                used_detections.add(best_match)
                old_pos = obj_data['last_position']
                new_pos = current_positions[best_match]
                # 移動距離の計算（表示用なので平方根を計算）
                move_distance = math.sqrt(self.calculate_distance_squared(old_pos, new_pos))

                matched_objects[obj_id] = {
                    'position': new_pos,
                    'last_position': new_pos,
                    'class_id': current_classes[best_match],
                    'score': current_scores[best_match],
                    'box': current_boxes[best_match],
                    'mask': current_masks[best_match],
                    'first_seen': obj_data['first_seen'],
                    'last_seen': self.frame_count,
                    'total_distance': obj_data['total_distance'] + move_distance,
                    'trail': obj_data['trail'] + [new_pos],
                    'lost_frames': 0
                }

                # 軌跡の長さを制限
                if len(matched_objects[obj_id]['trail']) > TRAIL_LENGTH:
                    matched_objects[obj_id]['trail'] = matched_objects[obj_id]['trail'][-TRAIL_LENGTH:]
            else:
                # マッチングしなかった場合、ロストフレーム数を増加
                obj_data['lost_frames'] += 1
                if obj_data['lost_frames'] <= MAX_LOST_FRAMES:
                    matched_objects[obj_id] = obj_data

        # 新しい検出をオブジェクトとして追加
        for i, pos in enumerate(current_positions):
            if i not in used_detections:
                matched_objects[self.next_id] = {
                    'position': pos,
                    'last_position': pos,
                    'class_id': current_classes[i],
                    'score': current_scores[i],
                    'box': current_boxes[i],
                    'mask': current_masks[i],
                    'first_seen': self.frame_count,
                    'last_seen': self.frame_count,
                    'total_distance': 0.0,
                    'trail': [pos],
                    'lost_frames': 0
                }
                self.next_id += 1

        self.tracked_objects = matched_objects
        return self.tracked_objects


def video_processing(frame, model, tracker, results_log, last_print_time):
    # YOLOEによる推論実行
    results = model(frame, conf=CONF_THRESHOLD, verbose=False)

    # 物体追跡の更新
    tracked_objects = tracker.update(results)

    current_time = time.time()
    processed_frame = frame.copy()

    # 追跡結果の描画と情報収集
    detection_info = []

    for obj_id, obj_data in tracked_objects.items():
        if obj_data['lost_frames'] == 0:  # 現在フレームで検出されている物体のみ
            class_name = model.names[obj_data['class_id']] if obj_data['class_id'] < len(model.names) else f'class_{obj_data["class_id"]}'
            score = obj_data['score']
            duration = obj_data['last_seen'] - obj_data['first_seen'] + 1
            total_distance = obj_data['total_distance']

            # セグメンテーションマスクの描画
            if obj_data['mask'] is not None:
                mask_resized = cv2.resize(obj_data['mask'], (frame.shape[1], frame.shape[0]))
                mask_colored = np.zeros_like(frame)
                # オブジェクトIDに基づいた一意な色生成
                color_seed = obj_id * 50
                color = (
                    (color_seed * 7) % 255,
                    (color_seed * 13) % 255,
                    (color_seed * 19) % 255
                )
                mask_colored[mask_resized > 0.5] = color
                processed_frame = cv2.addWeighted(processed_frame, 1, mask_colored, 0.3, 0)
            else:
                # マスクがない場合のデフォルト色
                color = (0, 255, 0)

            # バウンディングボックス描画
            box = obj_data['box']
            x1, y1, x2, y2 = map(int, box)
            cv2.rectangle(processed_frame, (x1, y1), (x2, y2), color, 2)

            # ID、クラス、スコア、滞在時間の表示
            label = f'ID:{obj_id} {class_name}: {score:.2f} ({duration}f)'
            cv2.putText(processed_frame, label, (x1, y1-10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            # 軌跡の描画
            trail = obj_data['trail']
            if len(trail) > 1:
                for i in range(1, len(trail)):
                    pt1 = (int(trail[i-1][0]), int(trail[i-1][1]))
                    pt2 = (int(trail[i][0]), int(trail[i][1]))
                    cv2.line(processed_frame, pt1, pt2, color, 2)

            # 検出情報の記録
            detection_info.append(f'ID:{obj_id} {class_name}: {score:.2f} (滞在:{duration}f, 移動:{total_distance:.1f}px)')

    # 1秒間隔での出力
    if current_time - last_print_time[0] >= PRINT_INTERVAL:
        if detection_info:
            output = f'[{time.strftime("%H:%M:%S")}] 追跡中: {", ".join(detection_info)}'
        else:
            output = f'[{time.strftime("%H:%M:%S")}] 追跡中: No objects'
        print(output)
        results_log.append(output)
        last_print_time[0] = current_time

    return processed_frame


# プログラム概要表示
print('=== YOLOEオブジェクト検出・追跡プログラム ===')
print('概要: リアルタイムでオブジェクトを検出・追跡し、セグメンテーションマスクと軌跡で表示します')
print('機能: YOLOEによるオープンボキャブラリ物体検出・セグメンテーション・追跡')
print('追跡情報: 物体ID、滞在時間、移動距離、軌跡を表示・記録')
print('操作: qキーで終了')
print()

# YOLOEモデル設定（デフォルト：s、変更可能：s, m, l）
print('モデル選択:')
print('1. yoloe-11s-seg.pt - 軽量')
print('2. yoloe-11m-seg.pt - バランス型')
print('3. yoloe-11l-seg.pt - 精度重視')
choice = input('選択[1-3]: ')
sizes = {'1': 's', '2': 'm', '3': 'l'}
MODEL_SIZE = sizes.get(choice, 's')
MODEL_NAME = f'yoloe-11{MODEL_SIZE}-seg.pt'

# GPU/CPU自動選択
if torch.cuda.is_available():
    device = 'cuda'
    print(f'GPU検出: {torch.cuda.get_device_name(0)}')
else:
    device = 'cpu'
    print('CPUモードで実行します')

# YOLOEモデル初期化
try:
    print(f'YOLOE{MODEL_SIZE}モデルを初期化中...')
    model = YOLOE(MODEL_NAME)

    # 重要: オープンボキャブラリ機能を有効にするためのクラス設定
    model.set_classes(CLASSES, model.get_text_pe(CLASSES))
    print(f'モデル読み込み完了。検出対象: {len(CLASSES)}クラス')
except Exception as e:
    print(f'モデルの読み込みに失敗しました: {e}')
    exit()

# 物体追跡器の初期化
tracker = ObjectTracker()
results_log = []
last_print_time = [time.time()]

print('0: 動画ファイル')
print('1: カメラ')
print('2: サンプル動画')

choice = input('選択: ')
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)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
elif choice == '2':
    # サンプル動画ダウンロード・処理
    url = SAMPLE_VIDEO_URL
    filename = SAMPLE_VIDEO_FILE
    try:
        urllib.request.urlretrieve(url, filename)
        temp_file = filename
        cap = cv2.VideoCapture(filename)
    except Exception as e:
        print(f'動画のダウンロードに失敗しました: {url}')
        print(f'エラー: {e}')
        exit()
else:
    print('無効な選択です')
    exit()

# メイン処理
try:
    while True:
        cap.grab()
        ret, frame = cap.retrieve()
        if not ret:
            break

        processed_frame = video_processing(frame, model, tracker, results_log, last_print_time)
        cv2.imshow(WINDOW_NAME, processed_frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
finally:
    cap.release()
    cv2.destroyAllWindows()
    if temp_file:
        os.remove(temp_file)

    # 結果保存（追跡統計情報を含む）
    with open(RESULT_FILE, 'w', encoding='utf-8') as f:
        f.write('=== YOLOEオブジェクト検出・追跡結果 ===\n')
        f.write(f'使用デバイス: {device.upper()}\n')
        f.write(f'総フレーム数: {tracker.frame_count}\n')
        f.write(f'追跡したオブジェクト数: {tracker.next_id - 1}\n')
        f.write('\n=== 追跡統計 ===\n')

        # 各オブジェクトの詳細統計
        for obj_id, obj_data in tracker.tracked_objects.items():
            class_name = model.names[obj_data['class_id']] if obj_data['class_id'] < len(model.names) else f'class_{obj_data["class_id"]}'
            duration = obj_data['last_seen'] - obj_data['first_seen'] + 1
            f.write(f'ID:{obj_id} {class_name} - 滞在:{duration}フレーム, 移動距離:{obj_data["total_distance"]:.1f}px\n')

        f.write('\n=== 時系列ログ ===\n')
        for log in results_log:
            f.write(log + '\n')
    print('result.txtに保存')