はじめに

物体検出のYOLOシリーズの最新版であるYOLOXを試してみました。
データセットにTACOデータセットを使用し、ゴミ検出AIを作るまでの作業を備忘録として残します。

環境
OS: Ubuntu18.04LTS
CPU: Intel® Core™ i7-8700 CPU @ 3.20GHz × 12
GPU: GeForce RTX2080
Python: Python3.6.5(Anacondaで作成した仮想環境)
scipy:0.19.0
numpy:1.16.0
keras:2.2.4
matplotlib:2.0.0
opencv:4.2.0
tensorflow-gpu:1.12.0
tqdm:4.11.2
pillow:8.1.0
h5py:2.8.0

YOLOXのダウンロード

個人的にkerasに馴染みがあるので、こちらからkeras版YOLOXをcloneします。

git clone https://github.com/bubbliiiing/yolox-keras

学習済みモデルのダウンロード

COCOで学習したモデルをダウンロードします。モデルのサイズ色々ありますが、今回はyolox_s.h5を使用します。

cd yolox-keras/model_data
wget https://github.com/bubbliiiing/yolox-keras/releases/download/v1.0/yolox_s.h5

wgetじゃなくても、YOLOXをcloneしてきたページ内のリンクから直接ダウンロードしてmodel_data内に置くのでもいいです。

TACOデータセットのダウンロード

こちらからyolox-kerasと同じフォルダ内にcloneし、download.pyを実行します。

git clone https://github.com/pedropro/TACO
cd TACO
python download.py

dataフォルダ内にbatxh_1~15のフォルダとその中にゴミの画像データがダウンロードされます。

データセットの形式の変換(COCO形式→VOC形式)

keras版YOLOXはVOC形式のデータセットに対応していますが、TACOはCOCO形式のデータセットになります。そこでこちらを使ってデータセットの形式を変換します。
上と同様にyolox-kerasと同じフォルダ内にcloneし、convert_coco_to_voc.pyを実行します。１つ目の引数(../TACO/data/annotations.json)がTACOのアノテーション情報が書かれたjsonファイル、２つ目の引数(./)が出力先、--folderオプションで出力フォルダの名前を指定しています。

git clone https://github.com/Kazuhito00/convert_coco_to_voc
cd convert_coco_to_voc
python convert_coco_to_voc.py ../TACO/data/annotations.json ./ --folder VOCTACO

学習の準備

生成されたVOCTACO/JPEGImages内に先ほどダウンロードしたTACOの画像データセットをコピーします。(batch_*フォルダごと)

cp -r ../TACO/data/batch_* ./VOCTACO/JPEGImages/

VOCTACOフォルダをyolox-keras/VOCdevkitフォルダ内に移動します。

mv VOCTACO ../yolox-keras/VOCdevkit/

TACOのサイトのTaxonomyページ内にある表からCategory列をコピペした、taco_classes.txtを作成しyolo-keras/model_data内に保存します。

taco_classes.txt

Aluminium foil
Battery
Aluminium blister pack
Carded blister pack
Clear plastic bottle
Glass bottle
Other plastic bottle
Plastic bottle cap
Metal bottle cap
Broken glass
Aerosol
Drink can
Food can
Corrugated carton
Drink carton
Egg carton
Meal carton
Pizza box
Toilet tube
Other carton
Cigarette
Paper cup
Disposable plastic cup
Foam cup
Glass cup
Other plastic cup
Food waste
Glass jar
Plastic lid
Metal lid
Normal paper
Tissues
Wrapping paper
Magazine paper
Paper bag
Plastified paper bag
Garbage bag
Single-use carrier bag
Polypropylene bag
Plastic Film
Six pack rings
Crisp packet
Other plastic wrapper
Spread tub
Tupperware
Disposable food container
Foam food container
Other plastic container
Plastic glooves
Plastic utensils
Pop tab
Rope
Scrap metal
Shoe
Squeezable tube
Plastic straw
Paper straw
Styrofoam piece
Other plastic
Unlabeled litter

yolox-keras内にあるvoc_annotation.pyを修正して実行します。修正箇所と内容は以下になります。

voc_annotation.py

〜略〜

classes_path        = 'model_data/taco_classes.txt'

〜略〜

VOCdevkit_sets  = [('TACO', 'train'), ('TACO', 'val')]

〜略〜
        
if __name__ == "__main__":
    random.seed(0)
    if annotation_mode == 0 or annotation_mode == 1:
        print("Generate txt in ImageSets.")
        xmlfilepath     = os.path.join(VOCdevkit_path, 'VOCTACO/Annotations')
        saveBasePath    = os.path.join(VOCdevkit_path, 'VOCTACO/ImageSets/Main')
        
〜略〜

    if annotation_mode == 0 or annotation_mode == 2:
        print("Generate 2007_train.txt and 2007_val.txt for train.")
        for year, image_set in VOCdevkit_sets:
            image_ids = open(os.path.join(VOCdevkit_path, 'VOC%s/ImageSets/Main/%s.txt'%(year, image_set)), encoding='utf-8').read().strip().split()
            list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
            for image_id in image_ids:
                # xmlから画像のpathを取得する
                in_file = open(os.path.join(VOCdevkit_path, 'VOC%s/Annotations/%s.xml'%(year, image_id)), encoding='utf-8')
                tree = ET.parse(in_file)
                root = tree.getroot()
                image_path = root[1].text[:-4]
    
                list_file.write('%s/VOC%s/JPEGImages/%s'%(os.path.abspath(VOCdevkit_path), year, image_path))

〜略〜

cd yolox-keras
python voc_annotation.py

実行後、yolox-keras内にTACO_train.txt、TACO_val.txtが生成され、
yolox-keras/VOCdevkit/VOCTACO/ImageSets/Main内にtrain.txt、trainval.txt、val.txt、test.txtが生成されます。

学習

yolox-keras内にあるtrain.pyを修正し実行します。修正箇所と内容は以下になります。

train.py

〜略〜
    classes_path    = 'model_data/taco_classes.txt'
〜略〜
    train_annotation_path   = 'TACO_train.txt'
    val_annotation_path     = 'TACO_val.txt'
〜略〜

python train.py

学習パラメータはデフォルトのままで計350epochの学習をさせたところ、1日半ほどで終わりました。

推論(画像)

yolox-keras内にあるpredict.pyを修正し実行します。修正箇所と内容は以下になります。学習時にyolox-keras/logフォルダ内に学習モデルが保存されるのでそのモデルのパスと、カテゴリー情報が書かれたtaco_classes.txtのパスにします。

predict.py

〜略〜
if __name__ == "__main__":
    yolo = YOLO(model_path="./logs/best_epoch_weights.h5", classes_path="./model_data/taco_classes.txt")
〜略〜

python predict.py

実行するとInputファイルのパスを聞かれるので、TACOデータセットの内の1枚のパスを入力し
推論をしてみます。

結果

ペットボトルとそのキャップに関しては検出できていますが、その他の缶や包装ゴミは検出に失敗しています。
何種類かの画像で試してみましたが、どうやらペットボトルやキャップは検出率が高く、他のカテゴリに関しては今一つといった精度でした。学習回数が足りなかったり、データセットの偏りが原因かなと思います。

推論(Webカメラ)

yolox-keras内にあるpredict.pyを修正し実行します。修正箇所と内容は以下になります。

predict.py

〜略〜
if __name__ == "__main__":
    〜略〜
    mode = "video"
〜略〜

python predict.py

結果

不安定ですが、一応検出は出来ています。