【AutoGluon】【物体検出】Pascal VOC datasetを用いて顔検出を学習する

2021年12月20日に記事を更新しました。

はじめに

以前GluonCVを使って同様のことをしました。
touch-sp.hatenablog.com
touch-sp.hatenablog.com
今回はAutoGluonを使って行います。

学習データ

以前と全く同じものを使いました。学習データの準備は過去記事を参照して下さい。

Pythonスクリプト

from autogluon.vision import ObjectDetector

dataset_train = ObjectDetector.Dataset.from_voc('VOCdevkit/VOC2012', splits='train')

time_limit = 5*60*60  # 5 hour
detector = ObjectDetector()
hyperparameters = {'batch_size':4}
hyperparameter_tune_kwargs={'num_trials': 2}
detector.fit(dataset_train,
            time_limit = time_limit, 
            hyperparameters = hyperparameters,
            hyperparameter_tune_kwargs = hyperparameter_tune_kwargs)

detector.save('detector.ag')

たったこれだけですか？そう聞きたくなるほど短いスクリプトで学習ができました。

結果

学習時にタイムリミットを5時間に設定していますが実際にかかった時間は約15分でした。

Finished, total runtime is 903.68 s

結果を確認するためのスクリプトがこちら。
学習のためのスクリプトよりもこちらのほうが長いです。結果を図で示す関数が見つからなかったので自分で書いたため長くなってしまいました。

import numpy as np
from matplotlib import pyplot as plt

from mxnet import image
from gluoncv import utils

from autogluon.vision import ObjectDetector

detector = ObjectDetector.load('detector.ag')

img_file = utils.download( 'https://raw.githubusercontent.com/dmlc/web-data/master/gluoncv/segmentation/mhpv1_examples/1.jpg')
image_array = image.imread(img_file)

result = detector.predict(image_array)
selected_result = result.query('predict_score > 0.8')

class_ids , class_names = selected_result['predict_class'].factorize()

bounding_boxes = np.array([[x[i] for i in x.keys()] for x in selected_result['predict_rois']])

scores = np.array(selected_result['predict_score'])

utils.viz.plot_bbox(image_array, bounding_boxes, scores=scores,
                    labels=class_ids, class_names = class_names, absolute_coordinates=False)

plt.show()

追記

「SSD」「YOLO3」「CenterNet」の3種類が用意されています。学習するモデルを限定する場合には以下のようにします。

from autogluon.vision import ObjectDetector
from autogluon .core.space import Categorical

dataset_train = ObjectDetector.Dataset.from_voc('VOCdevkit/VOC2012', splits='train')

time_limit = 5*60*60  # 5 hour
detector = ObjectDetector()
hyperparameters = {
    'batch_size':4, 
    'transfer': Categorical('ssd_512_resnet50_v1_coco'),
    'epochs': 10,
    'early_stop_patience': 5}

hyperparameter_tune_kwargs={'num_trials': 3}
detector.fit(dataset_train,
            time_limit = time_limit, 
            hyperparameters = hyperparameters,
            hyperparameter_tune_kwargs = hyperparameter_tune_kwargs)

detector.save('detector.ag')

動作環境

Intel(R) Core(TM) i7-7700K
RAM 32G
GTX 1080

Ubuntu 20.04 on WSL2
Python3.8.10
mxnet-cu112==1.9.0
autogluon==0.3.2b20211220