import cv2
import mediapipe as mp
# Tasks APIの正しいインポートパス
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

# 1. クラスのセットアップ (正しいパスから取得)
BaseOptions = python.BaseOptions
FaceDetector = vision.FaceDetector
FaceDetectorOptions = vision.FaceDetectorOptions
VisionRunningMode = vision.RunningMode

# 2. 顔検出器の設定
options = FaceDetectorOptions(
    base_options=BaseOptions(model_asset_path='blaze_face_short_range.tflite'),
    running_mode=VisionRunningMode.VIDEO
)

# Webカメラのキャプチャを開始 (デフォルトカメラは0)
cap = cv2.VideoCapture(0)

# FaceDetectorのインスタンスを生成
with FaceDetector.create_from_options(options) as detector:
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            print("Webカメラからの映像を取得できませんでした。")
            break

        # OpenCVのBGR画像をRGBに変換
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # タイムスタンプ（ミリ秒）を取得
        frame_timestamp_ms = int(cap.get(cv2.CAP_PROP_POS_MSEC))

        # 3. MediaPipe用のImageオブジェクトを作成
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)

        # 4. 検出の実行
        detection_result = detector.detect_for_video(mp_image, frame_timestamp_ms)

        # 5. 検出結果の描画
        if detection_result.detections:
            for detection in detection_result.detections:
                bbox = detection.bounding_box

                # 座標はピクセル単位で取得
                start_point = (bbox.origin_x, bbox.origin_y)
                end_point = (bbox.origin_x + bbox.width, bbox.origin_y + bbox.height)

                # 顔を囲む枠を描画 (緑色, 太さ2)
                cv2.rectangle(frame, start_point, end_point, (0, 255, 0), 2)

                # スコア（信頼度）の描画
                if detection.categories:
                    score = detection.categories[0].score
                    cv2.putText(frame, f"Face: {score:.2f}", (bbox.origin_x, bbox.origin_y - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # 結果を画面に表示
        cv2.imshow('MediaPipe Tasks Face Detection', frame)

        # 'q' キーを押すとループを抜ける
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

cap.release()
cv2.destroyAllWindows()