簡略学習探知AIカメラ

raspberrypiを使ってYOLOとOpenCVを使った物体検出をしたい。しかし、自分の望む物体を検出させるには画像のラベル化、学習とやることがたくさんある。それを簡略化するシステムを作成した。

まず、画像学習用の画像のラベル化は専用アプリの使用や操作が面倒なので簡単にする必要がある。

次に、raspberrypiだけでは処理が重いため学習は難しい、そのため普通は外部のGPUを使って学習を行うgooglecolabで学習を行うが、コードを描く必要があったり制限時間などの制限があるため面倒、今回は少しでも簡略化したいためraspberrypiだけで学習を行う方法を試してみた。

piカメラで撮影し、ディスプレイで描写し指でラベルを行う範囲を決定しラベル化を行う。
ラベルデータと画像データをもとに学習を開始する。
YOLOとOpencvで画像処理をしてターゲットの周りに線を描き、ディスプレイで見る。

まずpipで必要な物をインストールしていく


pip install opencv-python-headless pandas ultralytics cvzone picamera2 tk


from ultralytics import YOLO
import cv2
import os
import time
from picamera2 import Picamera2
import tkinter as tk

class TouchKeyboard(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Touch Keyboard")
        
        # 全画面表示
        self.attributes('-fullscreen', True)
        self.label = tk.Label(self, text="学習対象名", font=("NotoSansJP", 24))
        self.label.grid(row=0, column=0, columnspan=10, pady=20)

        self.entry = tk.Entry(self, width=50, font=("NotoSansJP", 18))
        self.entry.grid(row=1, column=0, columnspan=10, pady=20)
        # レイアウトの設定
        layout = [
            '1 2 3 4 5 6 7 8 9 0',
            'Q W E R T Y U I O P',
            'A S D F G H J K L',
            'Z X C V B N M',
            'Space Clear Enter'
        ]
        
        # キーを作成する
        for row_index, row in enumerate(layout):
            for col_index, key in enumerate(row.split()):
                if key == "Space":
                    btn = tk.Button(self, text="Space", width=15, height=2, font=("NotoSansJP", 14), command=lambda: self.type_key(' '))
                    btn.grid(row=row_index+2, column=col_index, columnspan=5, padx=5, pady=5, sticky="nsew")
                elif key == "Clear":
                    btn = tk.Button(self, text="Clear", width=15, height=2, font=("NotoSansJP", 14), command=self.clear_entry)
                    btn.grid(row=row_index+2, column=col_index+4, columnspan=3, padx=5, pady=5, sticky="nsew")
                elif key == "Enter":
                    btn = tk.Button(self, text="Enter", width=15, height=2, font=("NotoSansJP", 14), command=self.exit_with_output)
                    btn.grid(row=row_index+2, column=col_index+6, columnspan=2, padx=5, pady=5, sticky="nsew")
                else:
                    btn = tk.Button(self, text=key, width=5, height=2, font=("NotoSansJP", 14), command=lambda key=key: self.type_key(key))
                    btn.grid(row=row_index+2, column=col_index, padx=5, pady=5, sticky="nsew")

    def type_key(self, key):
        # エントリーにキーを追加する
        self.entry.insert(tk.END, key)

    def clear_entry(self):
        # エントリーのテキストをクリアする
        self.entry.delete(0, tk.END)

    def exit_with_output(self):
        # Enterキーが押されたときの処理
        global class_name
        class_name = self.entry.get()
        #set_class_name(class_name)
        update_data_yaml(class_name)        
        self.destroy()  # アプリケーションを終了する
        
class TouchKeyboard2(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Touch Keyboard")
        
        # 全画面表示
        self.attributes('-fullscreen', True)
        self.label = tk.Label(self, text="学習データ名", font=("NotoSansJP", 24))
        self.label.grid(row=0, column=0, columnspan=10, pady=20)

        self.entry = tk.Entry(self, width=50, font=("NotoSansJP", 18))
        self.entry.grid(row=1, column=0, columnspan=10, pady=20)

        # レイアウトの設定
        layout = [
            '1 2 3 4 5 6 7 8 9 0',
            'Q W E R T Y U I O P',
            'A S D F G H J K L',
            'Z X C V B N M',
            'Space Clear Enter'
        ]
        
        # キーを作成する
        for row_index, row in enumerate(layout):
            for col_index, key in enumerate(row.split()):
                if key == "Space":
                    btn = tk.Button(self, text="Space", width=15, height=2, font=("NotoSansJP", 14), command=lambda: self.type_key(' '))
                    btn.grid(row=row_index+2, column=col_index, columnspan=5, padx=5, pady=5, sticky="nsew")
                elif key == "Clear":
                    btn = tk.Button(self, text="Clear", width=15, height=2, font=("NotoSansJP", 14), command=self.clear_entry)
                    btn.grid(row=row_index+2, column=col_index+4, columnspan=3, padx=5, pady=5, sticky="nsew")
                elif key == "Enter":
                    btn = tk.Button(self, text="Enter", width=15, height=2, font=("NotoSansJP", 14), command=self.exit_with_output)
                    btn.grid(row=row_index+2, column=col_index+6, columnspan=2, padx=5, pady=5, sticky="nsew")
                else:
                    btn = tk.Button(self, text=key, width=5, height=2, font=("NotoSansJP", 14), command=lambda key=key: self.type_key(key))
                    btn.grid(row=row_index+2, column=col_index, padx=5, pady=5, sticky="nsew")

    def type_key(self, key):
        # エントリーにキーを追加する
        self.entry.insert(tk.END, key)

    def clear_entry(self):
        # エントリーのテキストをクリアする
        self.entry.delete(0, tk.END)

    def exit_with_output(self):
        # Enterキーが押されたときの処理
        global class_name_main
        class_name_main = self.entry.get()
        #set_class_name(class_name_main)
        update_data_yaml_num(num,class_name_main)      
        self.destroy()  # アプリケーションを終了する
        
class TouchKeyboard1(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Touch Keyboard")
        
        # 全画面表示
        self.attributes('-fullscreen', True)
        self.label = tk.Label(self, text="学習対象数", font=("NotoSansJP", 24))
        self.label.grid(row=0, column=0, columnspan=10, pady=20)

        self.entry = tk.Entry(self, width=50, font=("NotoSansJP", 18))
        self.entry.grid(row=1, column=0, columnspan=10, pady=20)

        # レイアウトの設定
        layout = [
            '1 2 3 4 5 6 7 8 9 0',
            'Space Clear Enter'
        ]
        
        # キーを作成する
        for row_index, row in enumerate(layout):
            for col_index, key in enumerate(row.split()):
                if key == "Space":
                    btn = tk.Button(self, text="Space", width=15, height=2, font=("NotoSansJP", 14), command=lambda: self.type_key(' '))
                    btn.grid(row=row_index+2, column=col_index, columnspan=5, padx=5, pady=5, sticky="nsew")
                elif key == "Clear":
                    btn = tk.Button(self, text="Clear", width=15, height=2, font=("NotoSansJP", 14), command=self.clear_entry)
                    btn.grid(row=row_index+2, column=col_index+4, columnspan=3, padx=5, pady=5, sticky="nsew")
                elif key == "Enter":
                    btn = tk.Button(self, text="Enter", width=15, height=2, font=("NotoSansJP", 14), command=self.exit_with_output)
                    btn.grid(row=row_index+2, column=col_index+6, columnspan=2, padx=5, pady=5, sticky="nsew")
                else:
                    btn = tk.Button(self, text=key, width=5, height=2, font=("NotoSansJP", 14), command=lambda key=key: self.type_key(key))
                    btn.grid(row=row_index+2, column=col_index, padx=5, pady=5, sticky="nsew")

    def type_key(self, key):
        # エントリーにキーを追加する
        self.entry.insert(tk.END, key)

    def clear_entry(self):
        # エントリーのテキストをクリアする
        self.entry.delete(0, tk.END)

    def exit_with_output(self):
        # Enterキーが押されたときの処理
        global num
        num = self.entry.get()
        self.destroy()
        
        # アプリケーションを終了する
        
# YOLOフォーマットに変換する関数
def convert_to_yolo_format(bbox, img_width, img_height):
    x_start, y_start, x_end, y_end = bbox
    x_center = (x_start + x_end) / 2.0 / img_width
    y_center = (y_start + y_end) / 2.0 / img_height
    width = (x_end - x_start) / img_width
    height = (y_end - y_start) / img_height
    return x_center, y_center, width, height

# ラベルを保存する関数
def save_yolo_labels(bbox, img_width, img_height, label_path, class_id=0):
    try:
        label_dir = os.path.dirname(label_path)
        if not os.path.exists(label_dir):
            os.makedirs(label_dir)

        x_center, y_center, width, height = convert_to_yolo_format(bbox, img_width, img_height)
        with open(label_path, 'w') as f:
            f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
        print(f"ラベルが保存されました: {label_path}")
    except Exception as e:
        print(f"ラベル保存時にエラーが発生しました: {e}")

# クラス名を設定する関数
def set_class_name(class_name,class_name_main):
    base_paths = [
       データ格納フォルダーのパス
    ]
    
    for base_path in base_paths:
        class_dir = os.path.join(base_path, class_name_main)
        if not os.path.exists(class_dir):
            os.makedirs(class_dir)
        
        classes_txt_path = os.path.join(class_dir, 'classes.txt')
        with open(classes_txt_path , 'a') as f:
            f.write(class_name + '\n')
        print(f"クラス名が設定されました: {class_name_main} (保存先: {classes_txt_path})")

# data.yamlを更新する関数
def update_data_yaml_num(num,class_name):
    data_yaml_path = '/home/wada/rpi-bookworm-yolov8-main/data.yaml'
    content = f"""# data.yaml
train:データ格納のパス/{class_name}
val: データ格納のパス/{class_name}

nc: {num}
names:
"""
    with open(data_yaml_path, 'w') as f:
        f.write(content)
    print("data.yamlが更新されました。")
    
def update_data_yaml(class_name):
    data_yaml_path = 'データ格納のパス'
    content = f"""  {i}: '{class_name}'
"""

    with open(data_yaml_path, 'a') as f:
        f.write(content)
    print("data.yamlが更新されました。")

# マウスコールバック関数
def mouse_callback(event, x, y, flags, param):
    global flag,x_start, y_start, x_end, y_end, drawing, box_set, img_width, img_height
    if event == cv2.EVENT_LBUTTONDOWN:
        # 撮影ボタンの範囲
        capture_button = (520, 220, 800, 300)
        exit_button = (20, 420, 200, 480)
        reset_button = (20, 20, 200, 80)

        if capture_button[0] <= x <= capture_button[2] and capture_button[1] <= y <= capture_button[3]:
            if box_set:
                print("撮影ボタンがクリックされました。")
                save_image_and_label(frame, img_width, img_height)  # 画像のサイズを渡す
            else:
                print("バウンディングボックスを作成してください。")
        elif exit_button[0] <= x <= exit_button[2] and exit_button[1] <= y <= exit_button[3]:
            print("終了ボタンがクリックされました。")
#             cv2.destroyAllWindows()
#             picam2.stop()
            flag = 1
        elif reset_button[0] <= x <= reset_button[2] and reset_button[1] <= y <= reset_button[3]:
            print("リセットボタンがクリックされました。")
            reset_bounding_box()
        else:
            if not box_set:
                drawing = True
                x_start, y_start = x, y
                x_end, y_end = x, y
    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            x_end, y_end = x, y
    elif event == cv2.EVENT_LBUTTONUP:
        if drawing and not box_set:
            drawing = False
            x_end, y_end = x, y
            box_set = True
            print(f"バウンディングボックス: x_start={x_start}, y_start={y_start}, x_end={x_end}, y_end={y_end}")

# バウンディングボックスをリセットする関数
def reset_bounding_box():
    global x_start, y_start, x_end, y_end, box_set
    x_start, y_start, x_end, y_end = 0, 0, 0, 0
    box_set = False
    print("バウンディングボックスがリセットされました。")

# 画像とラベルを保存する関数
def save_image_and_label(frame, img_width, img_height):
    global file_number, class_name_main,class_name
    base_file_name = f'{class_name}_{file_number}'

    image_path_train = f'データ格納のパス/{class_name_main}/{base_file_name}.jpg'
    label_path_train = f'データ格納のパス/{class_name_main}/{base_file_name}.txt'
    image_path_val = f'データ格納のパス/{class_name_main}/{base_file_name}.jpg'
    label_path_val = f'データ格納のパス/{class_name_main}/{base_file_name}.txt'

    # 画像を保存
    cv2.imwrite(image_path_train, frame)
    cv2.imwrite(image_path_val, frame)

    bbox = (x_start, y_start, x_end, y_end)
    save_yolo_labels(bbox, img_width, img_height, label_path_train)
    save_yolo_labels(bbox, img_width, img_height, label_path_val)

    print(f"画像とラベルが保存されました: {image_path_train}, {label_path_train} および {image_path_val}, {label_path_val}")

    file_number += 1


# プログラムを終了する関数
def exit_program():
    cv2.destroyAllWindows()
    picam2.stop()
    picam2.close()
    global class_name_main
    # YOLOv8の学習を再開可能な形で開始
    print("YOLOv8の学習を開始します...")
    model = YOLO('yolov8n.pt')

    print("新規に学習を開始します...")
    model.train(
     data='/home/wada/rpi-bookworm-yolov8-main/data.yaml',
     epochs=30,
      imgsz=640,
      batch=8,
      device='cpu',

     )

    model.save(f"best_{class_name_main}.pt")
    print("学習が完了しました。")

    print("プログラムを終了します。")
    exit(0)

# メイン処理
if __name__ == '__main__':
    try:
        global drawing, box_set, x_start, y_start, x_end, y_end, file_number, class_name, img_width, img_height
        global class_name_main
        global class_name
        global flag
        drawing = False
        box_set = False
        x_start, y_start, x_end, y_end = 0, 0, 0, 0
        file_number = 0

        # タッチディスプレイでクラスnumを入力する
        app1 = TouchKeyboard1()
        app1.mainloop()
        #name
        app2 = TouchKeyboard2()
        app2.mainloop()

        # Picamera2を初期化
        picam2 = Picamera2()
        camera_config = picam2.create_preview_configuration(main={"format": "RGB888", "size": (640, 480)})
        picam2.configure(camera_config)
        picam2.start()

        cv2.namedWindow("Object Detection", cv2.WND_PROP_FULLSCREEN)
        cv2.setWindowProperty("Object Detection", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        cv2.setMouseCallback("Object Detection", mouse_callback)


        for i in range(int(num)):
            print (class_name_main)
            
            app = TouchKeyboard()
            app.mainloop()
            print (class_name)
            set_class_name(class_name,class_name_main)
            flag = 0
            while flag != 1:
                # フレームをキャプチャ
                frame = picam2.capture_array()
                img_height, img_width, _ = frame.shape

                # フレームを上下反転
                frame = cv2.flip(frame, 0)

                # バウンディングボックスが設定されている場合は描画
                if box_set:
                    cv2.rectangle(frame, (x_start, y_start), (x_end, y_end), (255, 0, 0), 2)

                # 撮影ボタンの描画
                cv2.rectangle(frame, (520, 220), (800, 300), (0, 255, 0), -1)  # 緑色のボタン
                cv2.putText(frame, "Capture", (530, 260), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

                # 終了ボタンの描画
                cv2.rectangle(frame, (20, 420), (200, 480), (0, 0, 255), -1)  # 赤色のボタン
                cv2.putText(frame, "Exit", (50, 450), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

                # リセットボタンの描画
                cv2.rectangle(frame, (20, 20), (200, 80), (255, 255, 0), -1)  # 青色のボタン
                cv2.putText(frame, "Reset", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)

                # フレームを表示
                cv2.imshow("Object Detection", frame)

            # 'q'キーで終了
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    exit_program()
        
    
    finally:
        exit_program()


import cv2
from picamera2 import Picamera2
import pandas as pd
from ultralytics import YOLO
import cvzone
import os
import tkinter as tk
# 学習データフォルダーのパス
data_folder = 

# フォルダー内のファイルリストを取得
model_files = [f for f in os.listdir(data_folder) if f.endswith('.pt')]
model_files.sort()  # ファイル名をソート

# 選択されたモデル
selected_model = None
picam2 = None  # グローバルカメラオブジェクト

def select_model(model_name):
    global selected_model
    selected_model = model_name
    print(f"選択されたモデル: {selected_model}")
    model_window.destroy()  # モデル選択ウィンドウを閉じる
    restart_camera()  # カメラの再初期化を行う

def start_monitoring():
    open_model_selection()  # モデル選択ウィンドウを開く

def open_model_selection():
    global model_window
    # モデル選択ウィンドウ
    model_window = tk.Tk()
    model_window.title("モデル選択")
    
    # フルスクリーン設定
    model_window.attributes('-fullscreen', True)

    # モデル名をボタンとして表示
    for model in model_files:
        button = tk.Button(model_window, text=model, font=("Arial", 20), command=lambda m=model: select_model(m))
        button.pack(pady=20)

    # 戻るボタンを追加
    back_button = tk.Button(model_window, text="戻る", font=("Arial", 20), command=lambda: [model_window.destroy(), open_combined_screen()])
    back_button.pack(pady=20)

    model_window.protocol("WM_DELETE_WINDOW", model_window.quit)  # ウィンドウが閉じられたときの処理
    model_window.mainloop()  # メインループを開始

def restart_camera():
    global picam2, model, class_list

    # カメラが既に存在する場合は解放
    if picam2:
        picam2.stop()

    # カメラ設定
    picam2 = Picamera2()
    picam2.preview_configuration.main.size = (640, 480)
    picam2.preview_configuration.main.format = "RGB888"
    picam2.preview_configuration.align()
    picam2.configure("preview")
    picam2.start()

    # モデルの読み込み
    model = YOLO(os.path.join(data_folder, selected_model))  # 選択されたモデルを指定

    # クラス名の読み込み
    with open("coco1.txt", "r") as my_file:
        data = my_file.read()
    class_list = data.split("\n")

    # カメラのオブジェクト検出を開始
    detect_objects()

def detect_objects():
    count = 0
    # 複数の特定クラス名を指定
    target_class_names = ["person", "cat", "dog"]

    # ウィンドウをフルスクリーンに設定
    cv2.namedWindow("Camera", cv2.WND_PROP_FULLSCREEN)
    cv2.setWindowProperty("Camera", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

    while True:
        # 画像のキャプチャ
        im = picam2.capture_array()

        count += 1
        if count % 3 != 0:  # 3フレームに1回処理
            continue

        # 画像を上下反転
        im = cv2.flip(im, -1)
        results = model.predict(im)  # オブジェクト検出

        # 検出結果をデータフレームに変換
        a = results[0].boxes.data
        px = pd.DataFrame(a).astype("float")

        # 検出されたバウンディングボックスを描画
        detected = False  # 特定のクラスが検出されたかのフラグ
        for index, row in px.iterrows():
            x1, y1, x2, y2, _, d = row
            x1, y1, x2, y2, d = int(x1), int(y1), int(x2), int(y2), int(d)
            c = class_list[d]

            cv2.rectangle(im, (x1, y1), (x2, y2), (0, 0, 255), 2)
            cvzone.putTextRect(im, f'{c}', (x1, y1), 1, 1)

            # 複数の特定クラス名が検出された場合
            if c in target_class_names:
                detected = True

        # 特定のクラスが検出された場合の処理
        if detected:
            cvzone.putTextRect(im, "検出された", (50, 50), 2, 1, (0, 255, 0))

        # 画像を表示
        cv2.imshow("Camera", im)

        # マウスクリックイベントを待機
        if cv2.waitKey(1) & 0xFF == ord('b'):  # 'b'キーが押された場合
            picam2.stop()  # カメラを停止
            cv2.destroyAllWindows()  # ウィンドウを閉じる
            open_model_selection()  # モデル選択画面を再表示
            break

        # マウスイベントのチェック
        if cv2.getWindowProperty("Camera", cv2.WND_PROP_VISIBLE) >= 1:
            cv2.setMouseCallback("Camera", on_mouse_click)

        # 'q'が押されたら終了
        if cv2.waitKey(1) == ord('q'):
            break

    # 終了処理
    cv2.destroyAllWindows()

def on_mouse_click(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:  # 左ボタンがクリックされた場合
        picam2.stop()  # カメラを停止
        cv2.destroyAllWindows()  # ウィンドウを閉じる
        print("プログラムが終了しました")  # 終了メッセージ
        exit()  # プログラムを終了
        
if __name__ == "__main__":
    start_monitoring()  # 初期画面を表示