1. はじめに

オープンソースの物体検出アルゴリズムであるYOLOを用いて、料理画像から料理領域の抽出を行うモデルを作成したいと思います。
学習データとして、料理画像をインプットにして、料理領域のバウンディングボックスをラベルとして用います。
また、実行環境としてAWS SagemakerのNotebookインスタンス上のJupyter Notebookを使用しました。

2. YOLOとは

YOLOとはオープンソースのリアルタイム物体検出アルゴリズムで、物体の検出と識別を同時に行うことにより従来の物体検出アルゴリズムと比較して処理が軽いという特徴があります。また、汎化性能も高いです。
YOLOを使って物体検出をしてみた結果の例が下図になります。
物体同士が重なりあってたり、遠くにある場合でもしっかりと物体の検出と領域の抽出ができていることが分かります。

現在バージョン5まで発表されており、今回は最新のv5を使用したいと思います。
公式のサイトはこちらです。

3. 学習データ

学習データとして電気通信大学柳井研究室が公開している料理データセットであるUECFoodPixCimpleteを用いたいと思います。
学習用9000、評価用1000の計10000のデータセットであり、1つの料理画像に料理領域のセグメンテーションマスクがセットになっています。
一例としては以下です。

セグメンテーションマスクのRGBのRの値によってその領域が102クラスのうちどのクラスに属するのかが分かるようになっています。
ただ、Yoloの入力データにはバウンディングボックスが必要なので、セグメンテーションマスクの情報からバウンディングボックスの情報を抽出する前処理が必要になります。

4. やってみる

それでは実際に学習を行ってみたいと思います。

4.1 環境構築

今回はSagemakerのNotebookインスタンスを用いてYolo v5の構築・学習を行いたいと思います。インスタンスはGPUが好ましいですが、それほどハイスペックである必要はありません。
しかし、公式のサイトにあるように、Pythonのバージョン3.8以上が必要ですが、現在Notebookインスタンスにデフォルトで用意されているカーネルはPython3.6までしかありませんので、自分でPython3.8の環境を作る必要があります。
以前書いたこちらの記事を参考にPython3.8の環境を作ります。（手順は省略します）

4.2 Yolo v5のインストール・動作確認

まずはYolo v5のリポジトリをクローンします。ノートブック上で次のコマンドを打ちます。

%%sh
git clone https://github.com/ultralytics/yolov5  # clone repo

次に必要なモジュールのインストールを行います。

pip install -qr yolov5/requirements.txt

さらにセットアップを行います。

import torch
from IPython.display import Image, clear_output  # to display images

clear_output()
print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

ここまで来ると準備OKなので、サンプルのモデル・データで動作確認をします。
次のコマンドを打ってYolo v5で物体検知を起動します。

%%sh
cd yolov5
source activate /home/ec2-user/SageMaker/kernels/myenv
python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/

結果を見るために下のコマンドを打つと、サンプルの画像に対して物体検出を行った結果が出力されます。

Image(filename='yolov5/runs/detect/exp/zidane.jpg', width=600)

4.3 学習データの用意

次に今回の学習に用いる料理画像データとそのバウンディングボックスを用意します。
まずはUECFoodPixCimpleteをダウンロードします。

%%sh
wget http://mm.cs.uec.ac.jp/uecfoodpix/UECFOODPIXCOMPLETE.tar

次にこれを展開します。

# 展開
import tarfile

with tarfile.open('UECFOODPIXCOMPLETE.tar', 'r:tar') as tr:
    tr.extractall(path='.')

展開した際はディレクトリ構造がこんな感じになってます。

|--data
|  |--category.txt
|  |--test1000.txt
|  |--train9000.txt
|  |--UECFoodPIXCOMPLETE
|  |  |--test
|  |  |  |--img
|  |  |  |--mask
|  |  |--train
|  |  |  |--img
|  |  |  |--mask

ただ、Yoloの学習用にディレクトリ構造を次のように変えたいと思います。

|--data
|  |--category.txt
|  |--test1000.txt
|  |--train9000.txt
|  |--UECFoodPIXCOMPLETE
|  |  |--images
|  |  |  |--train
|  |  |  |--valid
|  |  |--labels
|  |  |  |--train
|  |  |  |--valid

次のコマンド群を叩きます。

%%sh
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/train
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/valid
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/labels
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/labels/train
mkdir ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/labels/valid
mv ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/train/img/* ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/train
mv ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/test/img/* ./UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/valid

labels以下にはYolo形式のバウンディングボックス情報を記載します。

Yolo形式のバウンディングボックスの書き方は下図の例のように1つの画像に対して、1オブジェクトを1行に、クラス、中央のx座標、中央のy座標、幅、高さを列にして表記します。詳しくは公式を参考してください。

セグメンテーションマスクからYolo形式のバウンディングボックスに変換する処理は次のコードです。
正直もっと効率良いやり方があると思ってます。。。
まずは学習用データ

import cv2  #OpenCVをインポート
import numpy as np #numpyをインポート
import glob
import os

# セグメンテーションマスク(pngファイル)のディレクトリ
files = glob.glob("UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/train/mask/*")
# yolo形式のテキストファイルを配置するディレクトリ
directory_yolo_txt = "UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/labels/train/"
# 料理画像(jpgファイル)のディレクトリ
jpg_files = "UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/train/"
for file in files: 
    img_png = cv2.imread(file) #画像の読み込み
    img_png = cv2.cvtColor(img_png, cv2.COLOR_BGR2RGB) #色配置の変換 BGR→RGB
    img_array = np.asarray(img_png) #numpyで扱える配列をつくる
    # 画像に含まれるラベルを保持するリスト
    labels = []
    # 各ラベルの右端の座標を保持するリスト
    right_x = []
    # 各ラベルの左端の座標を保持するリスト
    left_x = []
    # 各ラベルの上端の座標を保持するリスト
    upper_y = []
    # 各ラベルの下端の座標を保持するリスト
    lower_y = []
    # 各ラベルの中央のx座標(0~1の範囲に正規化)を保持するリスト
    cent_x = []
    # 各ラベルの中央のy座標(0~1の範囲に正規化)を保持するリスト
    cent_y = []
    # 各ラベルの高さ(0~1の範囲に正規化)を保持するリスト
    heights = []
    # 各ラベルの幅(0~1の範囲に正規化)を保持するリスト
    widths = []
    height, width = img_png.shape[:2]
    img_jpg = cv2.imread(jpg_files + file.split("/")[-1].split(".")[0] + ".jpg") #画像の読み込み
    height_jpg, width_jpg = img_jpg.shape[:2]
    # もしセグメンテーションマスクと料理画像の形が違う場合は飛ばす
    if height_jpg != height or width_jpg != width:
        continue
    for y,i in enumerate(img_array):
        for x,j in enumerate(i):
            if j[0] != 0:
                if j[0] not in labels:
                    labels.append(j[0])
                    right_x.append(x)
                    left_x.append(x)
                    upper_y.append(y)
                    lst_y.append(y)
                else:
                    if right_x[labels.index(j[0])] < x:
                        right_x[labels.index(j[0])] = x
                    if left_x[labels.index(j[0])] > x:
                        left_x[labels.index(j[0])] = x
                    if upper_y[labels.index(j[0])] < y:
                        upper_y[labels.index(j[0])] = y
    for i in range(0, len(labels)):
        cent_x.append((right_x[i] + left_x[i]) / 2 / width)
        cent_y.append((upper_y[i] + lower_y[i]) / 2 / height)
        widths.append((right_x[i] - left_x[i]) / width)
        heights.append((upper_y[i] - lower_y[i]) / height)
    # yolo形式でファイルを出力
    with open(directory_yolo_txt + file.split("/")[-1].split(".")[0] + ".txt" , "w", encoding="utf-8") as f:
        for i in range(0, len(labels)):
            strings = str(labels[i] -1) + " " + str(cent_x[i]) + " " + str(cent_y[i]) + " " + str(widths[i]) + " " + str(heights[i])
            f.write(strings)
            f.write("\n")

次に検証用データ

import cv2  #OpenCVをインポート
import numpy as np #numpyをインポート
import glob
import os

# セグメンテーションマスク(pngファイル)のディレクトリ
files = glob.glob("UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/test/mask/*")
# yolo形式のテキストファイルを配置するディレクトリ
directory_yolo_txt = "UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/labels/valid/"
# 料理画像(jpgファイル)のディレクトリ
jpg_files = "UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/valid/"
for file in files: 
    img_png = cv2.imread(file) #画像の読み込み
    img_png = cv2.cvtColor(img_png, cv2.COLOR_BGR2RGB) #色配置の変換 BGR→RGB
    img_array = np.asarray(img_png) #numpyで扱える配列をつくる
    # 画像に含まれるラベルを保持するリスト
    labels = []
    # 各ラベルの右端の座標を保持するリスト
    right_x = []
    # 各ラベルの左端の座標を保持するリスト
    left_x = []
    # 各ラベルの上端の座標を保持するリスト
    upper_y = []
    # 各ラベルの下端の座標を保持するリスト
    lower_y = []
    # 各ラベルの中央のx座標(0~1の範囲に正規化)を保持するリスト
    cent_x = []
    # 各ラベルの中央のy座標(0~1の範囲に正規化)を保持するリスト
    cent_y = []
    # 各ラベルの高さ(0~1の範囲に正規化)を保持するリスト
    heights = []
    # 各ラベルの幅(0~1の範囲に正規化)を保持するリスト
    widths = []
    height, width = img_png.shape[:2]
    img_jpg = cv2.imread(jpg_files + file.split("/")[-1].split(".")[0] + ".jpg") #画像の読み込み
    height_jpg, width_jpg = img_jpg.shape[:2]
    # もしセグメンテーションマスクと料理画像の形が違う場合は飛ばす
    if height_jpg != height or width_jpg != width:
        continue
    for y,i in enumerate(img_array):
        for x,j in enumerate(i):
            if j[0] != 0:
                if j[0] not in labels:
                    labels.append(j[0])
                    right_x.append(x)
                    left_x.append(x)
                    upper_y.append(y)
                    lst_y.append(y)
                else:
                    if right_x[labels.index(j[0])] < x:
                        right_x[labels.index(j[0])] = x
                    if left_x[labels.index(j[0])] > x:
                        left_x[labels.index(j[0])] = x
                    if upper_y[labels.index(j[0])] < y:
                        upper_y[labels.index(j[0])] = y
    for i in range(0, len(labels)):
        cent_x.append((right_x[i] + left_x[i]) / 2 / width)
        cent_y.append((upper_y[i] + lower_y[i]) / 2 / height)
        widths.append((right_x[i] - left_x[i]) / width)
        heights.append((upper_y[i] - lower_y[i]) / height)
    # yolo形式でファイルを出力
    with open(directory_yolo_txt + file.split("/")[-1].split(".")[0] + ".txt" , "w", encoding="utf-8") as f:
        for i in range(0, len(labels)):
            strings = str(labels[i] -1) + " " + str(cent_x[i]) + " " + str(cent_y[i]) + " " + str(widths[i]) + " " + str(heights[i])
            f.write(strings)
            f.write("\n")

4.4 学習

これで学習データの情報がそろったので、学習に進めます。
ただ、その前にラベルの情報や学習データの配置場所を記載したdataset.yamlというファイルを作成します。これをyolov5フォルダの直下に配置します。
中身は以下です。

dataset.yaml

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/train/
val: ../UECFOODPIXCOMPLETE/data/UECFoodPIXCOMPLETE/images/valid/

# number of classes
nc: 102

# class names
names: ['rice', 'eels on rice', 'pilaf', 'chicken-n-egg on rice', 'pork cutlet on rice', 'beef curry', 'sushi', 'chicken rice', 'fried rice', 'tempura bowl', 'bibimbap', 'toast', 'croissant', 'roll bread', 'raisin bread', 'chip butty', 'hamburger', 'pizza', 'sandwiches', 'udon noodle','tempura udon',
        'soba noodle', 'ramen noodle', 'beef noodle', 'tensin noodle', 'fried noodle', 'spaghetti', 'Japanese-style pancake', 'takoyaki', 'gratin',
        'sauteed vegetables', 'croquette', 'grilled eggplant', 'sauteed spinach', 'vegetable tempura', 'miso soup', 'potage', 'sausage',
        'oden', 'omelet', 'ganmodoki', 'jiaozi', 'stew', 'teriyaki grilled fish', 'fried fish', 'grilled salmon', 'salmon meuniere', 'sashimi',
        'grilled pacific saury ', 'sukiyaki', 'sweet and sour pork', 'lightly roasted fish', 'steamed egg hotchpotch', 'tempura', 'fried chicken',
        'sirloin cutlet', 'nanbanzuke', 'boiled fish','seasoned beef with potatoes', 'hambarg steak', 'beef steak', 'dried fish', 'ginger pork saute',
        'spicy chili-flavored tofu', 'yakitori', 'cabbage roll', 'rolled omelet', 'egg sunny-side up', 'fermented soybeans', 'cold tofu', 'egg roll',
        'chilled noodle', 'stir-fried beef and peppers', 'simmered pork', 'boiled chicken and vegetables', 'sashimi bowl', 'sushi bowl',
        'fish-shaped pancake with bean jam', 'shrimp with chill source', 'roast chicken', 'steamed meat dumpling', 'omelet with fried rice',
        'cutlet curry', 'spaghetti meat sauce', 'fried shrimp', 'potato salad', 'green salad', 'macaroni salad', 'Japanese tofu and vegetable chowder',
        'pork miso soup', 'chinese soup','beef bowl', 'kinpira-style sauteed burdock', 'rice ball', 'pizza toast', 'dipping noodles', 'hot dog',
        'french fries', 'mixed rice', 'goya chanpuru', 'others', 'beverage']

いよいよ学習です。以下のコードを打ちます。
今回はエポック数は100にしました。また、モデルは一番軽量な5sを指定しました。

%%sh
source activate /home/ec2-user/SageMaker/kernels/myenv
cd yolov5
# --device 0でGPU --device CPUでCPU
python train.py --img 640 --batch 16 --epochs 100 --data Dataset.yaml --cfg models/yolov5s.yaml --weights '' --device 0

ml.p2.8xlargeでだいたい12時間くらい時間がかかりました。

学習の途中経過はこんな感じです。

4.5 出来上がったモデルの動作確認

学習したモデルで料理領域検出をした結果が次になります。
ちゃんと検出できてますね。

5. まとめ

Yoloを用いて料理領域の検出と識別を行った。
より精度を上げていくには料理のクラスを増やしたり、領域をより細かくラベリングしていく必要がありそう。