main.py

import requests
import json
import base64
import os

config = {
    "40.png": {"black", "white", "blue", "yellow"},
    "38.png": {"blue", "white"},
    "20.png": {"black","white", "yellow", "blue"},
    "30.png": {"yellow", "blue", "white"},
    "3.jpg": {"black", "white"}
    # "28.png": ["yellow", "blue", "red"],
    # "35.png": ["blue", "green", "white"],
    # "15.png": ["white", "red", "yellow"],
    # "7.png": ["green", "white", "black"],
    # "37.png": ["white", "black", "blue"],
    # "31.png": ["red", "white", "yellow"],
    # "21.png": ["white", "black", "yellow"],
    # "34.png": ["yellow", "blue", "red"],
    # "22.png": ["white", "yellow", "blue"],
    # "10.png": ["green", "red", "white"],
    # "27.png": ["white", "yellow", "red"],
    # "25.png": ["white", "red", "green"],
    # "12.png": ["green", "red", "blue"],
    # "29.png": ["red", "blue", "white"],
    # "6.png": ["blue", "green", "white"],
    # "2.jpg": ["yellow", "white", "blue"],
    # "17.png": ["yellow", "blue", "black"],
    # "33.png": ["blue", "white", "yellow"],
    # "16.png": ["white", "yellow", "blue"],
    # "41.jpg": ["blue", "green", "white"],
    # "1.jpeg": ["white", "yellow", "blue"],
    # "26.png": ["green", "white", "red"],
    # "11.png": ["green", "red", "white"],
    # "14.png": ["green", "blue", "red"],
    # "42.jpg": ["red", "blue", "white"],
    # "4.png": ["black", "blue", "yellow"],
    # "23.png": ["yellow", "blue", "black"],
    # "36.png": ["green", "white", "blue"],
    # "24.png": ["white", "green", "blue"],
    # "9.png": ["green", "white", "black"],
    # "5.png": ["red", "white", "blue"],
    # "19.png": ["white", "yellow", "blue"],
    # "13.png": ["yellow", "blue", "black"],
    # "18.png": ["white", "black", "yellow"],
    # "8.png": ["black", "blue", "white"],
    # "32.png": ["black", "blue", "yellow"],
    # "39.png": ["white", "green", "black"]
}


def run_rec(i):
    with open(i, 'rb') as fp:
        img_base64 = base64.b64encode(fp.read())

        response = requests.post(
            'http://localhost:8501/v1/models/classifier:predict',
            json={
                "instances": [
                    {"b64": img_base64.decode('utf-8')}
                ]
            },
        )

        # print(response.headers)
        # print(response.text)
        # print(json.dumps(response.json(), indent=2))
        return response.json()


if __name__ == '__main__':

    for image in os.listdir('test-plates'):
        if image in config:
            res = run_rec(f'test-plates/{image}')
            predictions = res['predictions']
            if len(predictions) > 1:
                raise Exception(f'{image} predictions > 0')
            colors = predictions[0]['color']
            if set(colors) == set(config[image]):
                print(f'{image} - OK')
            else:
                print(f'{image} - FAIL - {set(colors)} != {set(config[image])}')