yolov8の検出結果のDB格納

検出結果のDB格納

vim create_table.py

を作成

import sqlite3

def create_table():
    conn = sqlite3.connect('detections.db')
    cursor = conn.cursor()

    # テーブルを作成
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS detections (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            label TEXT NOT NULL,
            count INTEGER NOT NULL,
            timestamp TEXT NOT NULL
        )
    ''')

    conn.commit()
    conn.close()

if __name__ == '__main__':
    create_table()

これを実行し
DBを作成

次にDBへ保存するモジュールの作成

vim inventory_database_module.py

でファイルを作成

import sqlite3
from datetime import datetime

def save_detection_to_db(detections):
    conn = sqlite3.connect('detections.db')
    cursor = conn.cursor()

    # 現在の時刻を取得
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    # 検出結果をテーブルに挿入
    for label, count in detections.items():
        cursor.execute('''
            INSERT INTO detections (label, count, timestamp)
            VALUES (?, ?, ?)
        ''', (label, count, current_time))

    conn.commit()
    conn.close()

として保存

import argparse
import json
import cv2
from ultralytics import YOLO
from collections import defaultdict
from line_notify import send_line_notify  # インポートを追加
from datetime import datetime
from inventory_database_module import save_detection_to_db  # データベース保存用の関数をインポート

# コマンドライン引数の解析
parser = argparse.ArgumentParser(description="YOLOv8 Object Detection")
parser.add_argument('image_path', type=str, help='Path to the input image file')
args = parser.parse_args()

# ラベルマッピングファイルのパス
label_mapping_path = 'label_mapping.json'

# JSONファイルからクラスラベルのマッピングを読み込み
with open(label_mapping_path, 'r', encoding='utf-8') as f:
    label_mapping = json.load(f)

# YOLOv8モデルのロード
model = YOLO('inventory_model/best.pt')  # ここで適切なモデルを選択

# 画像のロード
image = cv2.imread(args.image_path)

# 画像の検出
results = model(image, save=True, conf=0.1, iou=0.5)

# 検出結果の取得
detections = results[0]  # 最初の結果を取得
classes = detections.boxes.cls

# 検出物体のカウント
object_counts = defaultdict(int)
for cls in classes:
    class_label = model.names[int(cls)]
    if class_label in label_mapping:
        label = label_mapping[class_label]
    else:
        label = class_label
    object_counts[label] += 1

# 検出結果のフィルタリング（1以下のもの）
filtered_object_counts = {label: count for label, count in object_counts.items() if count <= 1}

# フィルタリングされた検出結果のメッセージ生成
message_lines = [f'{label}: {count}個' for label, count in filtered_object_counts.items()]
message = '\n'.join(message_lines)

# 現在の時刻を取得
current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
message = f"{message}\n在庫チェックの時刻: {current_time}"

# 検出結果の表示
for line in message_lines:
    print(line)

# LINE Notifyにメッセージを送信（フィルタリングされた結果のみ）
if message_lines:
    send_line_notify(message)
    save_detection_to_db(filtered_object_counts)  # データベースに検出結果を保存
else:
    print("No objects with counts of 1 or less detected.")

というように
結果をDBに保存し
在庫チェックの時刻も送信するようにコード変更

なお実行した後に
DBの中身を見るには

vim view_detections.py

で

import sqlite3

def view_detections(db_path='detections.db'):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    # テーブルの内容を取得
    cursor.execute('SELECT * FROM detections')
    rows = cursor.fetchall()

    # カラム名を取得
    column_names = [description[0] for description in cursor.description]

    # 結果を表示
    print(f"{' | '.join(column_names)}")
    print("-" * 50)
    for row in rows:
        print(" | ".join(str(value) for value in row))

    conn.close()

if __name__ == '__main__':
    view_detections()

として保存

python view_detections.py

を実行すると

id | label | count | timestamp
--------------------------------------------------
1 | バスクリン | 1 | 2024-07-07 06:45:47
2 | バスクリン | 1 | 2024-07-07 06:50:42
3 | バスクリン | 1 | 2024-07-07 06:51:43

となり検出結果の確認ができる

認識精度が低いため
精度を0.1まで下げないと認識しないし
並べた時の複数の検出ができていない

とりあえず指定ディレクトリの画像から検出するようにコード変更する