chatgptで作成した raspberry pi 5 + usb camera の動画を YOLO9 と OpenAPI で処理する python プログラム

sample-with_usb_camera-3.py

import cv2
import openai
import os
import sys
import logging
from contextlib import contextmanager
from dotenv import load_dotenv
from ultralytics import YOLO

# 環境変数からAPIキーを読み込み
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# ultralyticsのログ出力を抑制
logging.getLogger('ultralytics').setLevel(logging.CRITICAL)

# 標準出力と標準エラー出力を一時的に無効化するコンテキストマネージャ
@contextmanager
def suppress_output():
    with open(os.devnull, 'w') as devnull:
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        sys.stdout = devnull
        sys.stderr = devnull
        try:
            yield
        finally:
            sys.stdout = old_stdout
            sys.stderr = old_stderr

# YOLOv9モデルをロード
model = YOLO("yolov9s.pt")

def detect_objects(frame):
    # 標準出力と標準エラー出力を抑制してYOLOv9で推論
    with suppress_output():
        results = model(frame)
    
    detections = []
    for result in results:
        for box in result.boxes:
            # バウンディングボックスの座標とクラス名を取得
            cls = int(box.cls[0])
            label = result.names[cls]
            bbox = box.xyxy[0].tolist()  # xyxy (x1, y1, x2, y2)形式
            detections.append((label, bbox))
    
    return detections

def draw_boxes(frame, detections):
    for label, bbox in detections:
        x1, y1, x2, y2 = map(int, bbox)
        # 赤枠でバウンディングボックスを描画
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
        # クラス名を枠の上に表示
        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

def generate_description(detections):
    labels = [label for label, _ in detections]
    prompt = f"以下の物体を検出しました: {', '.join(labels)}. これらの物体について説明してください。"
    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "あなたは物体認識の結果を説明するアシスタントです。"},
                {"role": "user", "content": prompt}
            ],
            max_tokens = 300
        )
        return response.choices[0].message['content'].strip()
        
    except Exception as e:
        print(f"OpenAI API エラー: {e}")
        return "説明生成に失敗しました。"

if __name__ == "__main__":
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("カメラにアクセスできません。カメラデバイス番号を確認してください。")
        exit()

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("フレームをキャプチャできませんでした。")
                break

            # オブジェクト検出
            detections = detect_objects(frame)

            # 検出結果に基づいて赤枠を描画
            draw_boxes(frame, detections)

            # フレームを表示
            cv2.imshow("USB Camera - YOLOv8 Object Detection", frame)
            key = cv2.waitKey(10) & 0xFF

            if key == ord(' '):
                print("SPACEキーが押されました。説明を生成します...")
                if detections:
                    description = generate_description(detections)
                    print("AIからの説明:", description)
                else:
                    print("オブジェクトが検出されませんでした。")
                print("-----------------")

            elif key == ord('q'):
                print("プログラムを終了します...")
                break

    finally:
        cap.release()
        cv2.destroyAllWindows()