Python code to train and test with Custom Vision of Cognitive Service

Premise

--Create a Cognitive Service Custom Vision Service resource in Azure and create a project --Select the project type to Classification --Prepare 5 or more images for learning per tag, and create and store a folder for each tag. --Prepare a separate image to test the model

Uploading images for learning and executing learning

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

#Set the following information in the environment variables in advance
# 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'])

#Create a list of tags
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"))

#Specify the path of the image corresponding to the tag
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')

#Create a list of images to upload
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]))

#Upload up to 64 items of maximum size
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)

#Training execution
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)

Test the created model with a test image

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

#Set the following information in the environment variables in advance
# 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'])

#Specify Iteration to test
iterations = trainer.get_iterations(project_id)
target_iteration = iterations[0]

print('target iteration:', target_iteration.name)

#Specify the file to use in the test
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)

Output image

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
...

Relation

Amplify images for machine learning with python-Qiita