前提
- AzureでCognitive ServiceのCustom Vision Serviceのリソースを作り、プロジェクトを作成しておく
- プロジェクトの種類はClassificationにする
- 学習用の画像を1タグあたり5枚以上用意し、タグごとにフォルダを作成して格納しておく
- モデルをテストするための画像も別に用意しておく
学習用画像のアップロードと学習実行
train.py
import os, glob, time
from PIL import Image
from azure.cognitiveservices.vision.customvision.training import CustomVisionTrainingClient
from azure.cognitiveservices.vision.customvision.training.models import ImageFileCreateEntry
# 事前に以下の情報を環境変数に設定しておく
# export CUSTOMVISION_ENDPOINT=<your API endpoint>
# export CUSTOMVISION_TRAINING_KEY=<your training key>
# export CUSTOMVISION_PROJECT_ID=<your project id>
project_id = os.environ['CUSTOMVISION_PROJECT_ID']
trainer = CustomVisionTrainingClient(os.environ['CUSTOMVISION_TRAINING_KEY'], endpoint=os.environ['CUSTOMVISION_ENDPOINT'])
# タグのリストを作成
tag_list = []
tag_list.append(trainer.create_tag(project_id, "tagA"))
tag_list.append(trainer.create_tag(project_id, "tagB"))
tag_list.append(trainer.create_tag(project_id, "tagC"))
# タグに対応する画像のパスを指定
path_list = []
path_list.append('train-images/train/tagA/*.JPG')
path_list.append('train-images/train/tagB/*.JPG')
path_list.append('train-images/train/tagC/*.JPG')
# アップロード対象の画像リストを作成
image_list = []
for i, tag in enumerate(tag_list):
for file_path in glob.glob(path_list[i]):
with open(file_path, "rb") as image_contents:
image_list.append(ImageFileCreateEntry(name=os.path.basename(file_path), contents=image_contents.read(), tag_ids=[tag.id]))
# 最大サイズの64件ずつアップロード
for i in range(0, len(image_list), 64):
print('uploading', i, '-', i+64, '/', len(image_list))
upload_result = trainer.create_images_from_files(project_id, images=image_list[i:i+64])
if not upload_result.is_batch_successful:
print("Image batch upload failed.")
for image in upload_result.images:
print("Image status: ", image.status)
exit(-1)
# トレーニング実行
print ("Training...")
iteration = trainer.train_project(project_id)
while (iteration.status != "Completed"):
iteration = trainer.get_iteration(project_id, iteration.id)
print ("Training status: " + iteration.status)
time.sleep(5)
作成したモデルをテスト用画像でテスト
quickTest.py
import os, glob
from PIL import Image
from azure.cognitiveservices.vision.customvision.training import CustomVisionTrainingClient
from azure.cognitiveservices.vision.customvision.training.models import ImageFileCreateEntry
def probability(prediction):
return prediction.probability
# 事前に以下の情報を環境変数に設定しておく
# export CUSTOMVISION_ENDPOINT=<your API endpoint>
# export CUSTOMVISION_TRAINING_KEY=<your training key>
# export CUSTOMVISION_PROJECT_ID=<your project id>
project_id = os.environ['CUSTOMVISION_PROJECT_ID']
trainer = CustomVisionTrainingClient(os.environ['CUSTOMVISION_TRAINING_KEY'], endpoint=os.environ['CUSTOMVISION_ENDPOINT'])
# テストするIterationを指定
iterations = trainer.get_iterations(project_id)
target_iteration = iterations[0]
print('target iteration:', target_iteration.name)
# テストで使用するファイルを指定
files = glob.glob('train-images/test/**/*.JPG', recursive=True)
for i, file in enumerate(files):
image = open(file, 'rb').read()
result = trainer.quick_test_image(project_id, image, iteration_id=target_iteration.id)
max_prediction = max(result.predictions, key=probability)
print(file, max_prediction.tag_name, max_prediction.probability)
出力イメージ
target iteration: Iteration 1
train-images/test/tagA/IMG_001.JPG tagA 0.9999913
train-images/test/tagA/IMG_002.JPG tagA 0.9999975
train-images/test/tagB/IMG_003.JPG tagB 0.9999996
...