初心者がデータセットを少しきれいにした

はじめに

サメの画像認識AIを作りたいが、データセットの中には人や船など関係ないものがあり、
学習時の妨げになる可能性が高いと思いました。
そこで物体検出に特化したモデルを使い、サメ以外の情報が含まれている画像を
省きました。そのやり方についてまとめます

使ったデータセット

kaggleのサメに関してのデータセットを使いました

使用したモデル

使用したモデルは以下です

SSD MobileNet V2は、物体検出に特化したディープラーニングモデルです。
このモデルは、MobileNet V2をベースにした軽量な構造を持ち、モバイルデバイスや
リソースが限られた環境でもリアルタイムで動作するように設計されています

ソースコード

実装したコードについては以下です。

trim.py

import tensorflow as tf
import cv2
import numpy as np
import glob
import os
import tensorflow_hub as hub

# モデルをロード
model = hub.load("https://www.kaggle.com/models/tensorflow/ssd-mobilenet-v2/TensorFlow2/fpnlite-320x320/1")

# ラベルマップ（COCOデータセットに基づく）
category_to_remove = [1, 9]  # 1: person, 9: boat

# 入力ディレクトリのパス
input_dir = r'C:\Users\PC_User\Desktop\SharkAI\sharks'

# 出力ディレクトリのパス
output_dir = r'C:\Users\PC_User\Desktop\SharkAI\sharksProcessed'

# 入力ディレクトリ内のすべての画像を取得（サブディレクトリも含む）
image_paths = glob.glob(os.path.join(input_dir, '**', '*.jpg'), recursive=True)
image_paths += glob.glob(os.path.join(input_dir, '**', '*.jpeg'), recursive=True)
image_paths += glob.glob(os.path.join(input_dir, '**', '*.png'), recursive=True)
image_paths = list(set(image_paths))  # 重複を削除


# サブディレクトリを含む出力ディレクトリを作成
os.makedirs(output_dir, exist_ok=True)

# 処理した画像と失敗した画像のカウント
processed_images = 0
failed_images = 0

# サメ以外の物体（人や船）を検出し、その画像をスキップする
for image_path in image_paths:
    try:
        # 画像が存在するか確認
        if not os.path.exists(image_path):
            print(f"ファイルが見つかりません: {image_path}")
            failed_images += 1
            continue
        
        # 画像を読み込む
        image = cv2.imread(image_path)
        
        # 画像が正しく読み込まれたか確認
        if image is None:
            print(f"画像の読み込みに失敗しました: {image_path}")
            failed_images += 1
            continue
            
        # 画像をRGBに変換
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # モデルを使って物体検出
        input_tensor = tf.convert_to_tensor(image_rgb)
        input_tensor = input_tensor[tf.newaxis,...]  # バッチサイズの次元を追加

        # 検出結果
        detections = model(input_tensor)

        # 検出結果から必要なデータを抽出
        detection_boxes = detections['detection_boxes'][0].numpy()  # バウンディングボックス
        detection_classes = detections['detection_classes'][0].numpy().astype(np.int32)  # クラスID
        detection_scores = detections['detection_scores'][0].numpy()  # 信頼度

        # 人や船を検出した画像をスキップする
        has_person_or_boat = False
        for i in range(len(detection_boxes)):
            if detection_scores[i] > 0.5:  # 信頼度が50%以上の検出結果を処理
                class_id = detection_classes[i]
                if class_id in category_to_remove:  # サメ以外の物体（人や船）を検出
                    has_person_or_boat = True
                    print(f"人や船を検出したためスキップ: {image_path}")
                    break

        if has_person_or_boat:
            continue  # 次の画像へ

        # 出力ディレクトリ内の対応するパスを作成
        relative_path = os.path.relpath(image_path, input_dir)  # 元のパスを相対パスとして取得
        output_image_path = os.path.join(output_dir, relative_path)  # 出力パスを作成

        # 出力ディレクトリに対応するサブディレクトリを作成
        os.makedirs(os.path.dirname(output_image_path), exist_ok=True)

        # 画像を保存
        cv2.imwrite(output_image_path, image)
        processed_images += 1
        print(f"処理して保存: {output_image_path}")
    
    except Exception as e:
        print(f"エラーが発生しました: {image_path}")
        print(f"エラー内容: {str(e)}")
        failed_images += 1
        continue

# 結果の表示
print(f"\n処理完了!")
print(f"処理した画像: {processed_images}")
print(f"失敗した画像: {failed_images}")
print(f"全体の画像数: {len(image_paths)}")

このコードは生成AIにプロンプトを投げ作成してもらいました。
少し、手直しは必要でしたが10分程度でこのコードを書けるのですごいですね

実行結果

実際に実行してみましょう。↓が結果です

処理完了!
処理した画像: 1417
失敗した画像: 1
全体の画像数: 1546

うまくいってそう...？、画像の確認すべてしたわけでないですが、いったん人や船が含まれているものは消えてそうでした

まとめ

初心者ながらデータセットを少しきれいにしてみました。
この方法が正しいかはわかりませんが、自分がしたかったことはできたので満足です