はじめに

YOLOv4で物体検知モデルを作成する過程で、Open-ImagesというGoogleが提供しているデータセットを使用したのですが、その際地味に躓いたのでやった事を書きました。

環境

Google Colab

手順

1.fiftyoneをインストール

! pip install 'fiftyone<0.15.0'

2022/10/26 時点
colabの環境的に0.15.0より新しいバージョンを入れると依存関係でエラーが起きるので、0.14系を入れています
jinja2とfriskがconflictしてるようです

2.データセットをダウンロード

Open-Images-v6から、データセットを読みこみます
今回は、以下の条件で行っています

学習用・教師用・検証用の全てをダウンロード
信号機と標識が含まれた画像
画像の上限枚数:100枚

import fiftyone.zoo as foz

dataset = foz.load_zoo_dataset(
   "open-images-v6",  # データセット提供元
   # split="train",  # train(学習用), validation(検証用), test(テスト用)
   label_types=["detections"],
   classes=["Traffic light", "Traffic sign"],  # 使用する物体
   max_samples=100,  # 画像の上限枚数
   only_matching=True,
)
# 名前を付けて
dataset.name = "open-images-v6-traffic-light-sign"
# 永続化する
dataset.persistent = True

因みに、読みこんだ画像は以下のように保存されています
学習用・教師用・検証用で分けられているのが確認できます

root/
  ┗ fiftyone/
      ┗ open-images-v6/
         ┣ train/
         ┃  ┣ labels/
         ┃  ┣ metadata/
         ┃  ┗ data/
         ┃       ┣ image01.jpg
         ┃       ┣ image01.txt
         ┃       ┣ image02.jpg
         ┃       ┣ image02.txt
         ┃               ︙
         ┣ test/
         ┃  ┣ labels/
         ┃  ┣ metadata/
         ┃  ┗ data/
         ┃       ┣ image03.jpg
         ┃       ┣ image03.txt
         ┃       ┣ image04.jpg
         ┃       ┣ image04.txt
         ┃                ︙
         ┗ validation/
             ┣ labels/
             ┣ metadata/
             ┗ data/
                  ┣ image05.jpg
                  ┣ image05.txt
                  ┣ image06.jpg
                  ┣ image06.txt
                         ︙

3.データセットの中身を確認

import fiftyone as fo
session = fo.launch_app(dataset)

こんな感じで表示されます

4.データセットをGoogleDrive上にexport

手順2でGoogleColab上に保存したデータセットをYOLOv4形式でエクスポート、ついでに自分のGoogleDrive上に保存します
(セッションが切れるとGoogleColab上に保存したデータセットは消えるため)
今回は、export先としてtraffic_light_signというディレクトリを作成・指定してます

# googledriveをマウント
from google.colab import drive
drive.mount('/content/drive')
classes = ["Traffic light", "Traffic sign"]
export_ds = dataset

# エクスポート
export_ds.export(
    export_dir=f"/content/drive/MyDrive/traffic_light_sign",  # googledrive上に保存する
    dataset_type=fo.types.YOLOv4Dataset,  # 形式 : YOLOv4
    classes=classes,
    label_field="detections",
)

5.学習用・教師用・検証用の画像パスをそれぞれ書き出す

YOLOv4では、全画像を1つのディレクトリに詰め込み、学習用画像・教師用画像の区別は画像パスが記述されたテキストファイルを基にしています

exportの際にそのテキストファイルも出力してほしい... のですが残念ながらしてくれないようです
ただし、手順2で触れた通り区別自体はしてくれていますので、train・validation・testディレクトリ以下に保存されている画像のパスを読みに行く形で書き出していきます

from glob import glob
import os

# 画像パスの取得
train_images = glob('/root/fiftyone/open-images-v6/train/data/*.jpg')  # 学習用
val_images = glob('/root/fiftyone/open-images-v6/validation/data/*.jpg')  # 検証用
test_images = glob('/root/fiftyone/open-images-v6/test/data/*.jpg')  # 教師用

# 学習用の画像のパスをtrain.txtに書き込む
with open('/content/drive/MyDrive/traffic_light_sign/train.txt', 'w') as f:
    for line in train_images:
        f.write('data/%s\n' % os.path.basename(line))  # 例) data/image01.jpg

# 検証用の画像のパスをval.txtに書き込む
with open('/content/drive/MyDrive/traffic_light_sign/val.txt', 'w') as f:
    for line in val_images:
        f.write('data/%s\n' % os.path.basename(line))

# 教師用の画像のパスをtest.txtに書き込む
with open('/content/drive/MyDrive/traffic_light_sign/test.txt', 'w') as f:
    for line in test_images:
        f.write('data/%s\n' % os.path.basename(line))

わざわざこんな事しなくても普通にシャッフルすれば良い気もしますが、せっかく分けられているのでそれを使うことにしました

処理が終了したら、ディレクトリ内を確認してみましょう
最終的にこんな感じのファイル群が揃います

traffic_light_sign/
              ┣ data/  
              ┃    ┣ image01.jpg
              ┃    ┣ image01.txt  # image01.jpgのアノテーションファイル
              ┃    ┣ image02.jpg
              ┃    ┣ image02.txt  # image02.jpgのアノテーションファイル
              ┃            ︙
              ┣ obj.names   # 物体名が記述されているファイル
              ┣ train.txt   # 学習用画像のパスが記述されているファイル
              ┣ test.txt    # 教師用画像のパスが記述されているファイル
              ┣ val.txt     # 検証用画像のパスが記述されているファイル(そこまでいらない)
              ┗ images.txt  # 全画像のパスが記述されているファイル(そこまでいらない)

参考