teach_classificationbox.py

"""
Script to teach classificationbox image classes as described in
https://blog.machinebox.io/how-anyone-can-build-a-machine-learning-image-classifier-from-photos-on-your-hard-drive-very-5c20c6f2764f
https://github.com/machinebox/toys/blob/master/imgclass/main.go
Run from within the root folder containing the folders of images.
"""
import os
import requests

IP = 'localhost'
PORT = '8080'
CLASSIFIER = 'classificationbox'
VALID_FILETYPES = ('.jpg', '.png', '.jpeg')

TEACH_URL = f"http://{IP}:{PORT}/{CLASSIFIER}/teach"
HEALTH_URL = f"http://{IP}:{PORT}/readyz"


def check_classifier_health():
    """Check that classifier is reachable"""
    try:
        response = requests.get(HEALTH_URL)
        if response.status_code == 200:
            print("{} health-check passed".format(CLASSIFIER))
            return True
        else:
            print("{} health-check failed".format(CLASSIFIER))
            print(response.status_code)
            return False
    except requests.exceptions.RequestException as exception:
        print("{} is unreachable".format(CLASSIFIER))
        print(exception)


def list_folders(directiory='.'):
    """Returns a list of folders in a dir, defaults to current dir.
    These are not full paths, just the folder."""
    folders = [dir for dir in os.listdir(directiory)
               if os.path.isdir(os.path.join(directiory, dir))
               and not dir.startswith(directiory)
               and not dir.startswith('.')]
    folders.sort(key=str.lower)
    print("Folders found:")
    print(folders)
    return folders


def teach_name_by_file(teach_url, name, file_path):
    """Teach facebox a single name using a single file."""
    file_name = file_path.split("/")[-1]
    file = {'file': open(file_path, 'rb')}
    data = {'name': name, "id": file_name}

    response = requests.post(teach_url, files=file, data=data)

    if response.status_code == 200:
        print("File:{} taught with name:{}".format(file_name, name))
        return True

    elif response.status_code == 400:
        print("Teaching of file:{} failed with message:{}".format(
            file_name, response.text))
        return False

    elif response.status_code == 404:
        print("Teaching of file:{} failed due to :{}".format(
            file_name, response.text))
        return False

def main():
    if check_classifier_health():
        for folder_name in list_folders():
            folder_path = os.path.join(os.getcwd(), folder_name)
            print(f"Entering folder {folder_path}")
            for file in os.listdir(folder_path):
                
                if file.endswith(VALID_FILETYPES):
                    file_path = os.path.join(folder_path, file)
                    print(file_path)
                    teach_name_by_file(teach_url=TEACH_URL,
                                       name=folder_name,
                                       file_path=file_path)


if __name__ == '__main__':
    main()