predict.py

'''
File: recognize.py
Project: BulgarianPatternsRecognizer
Author: Milko Videv (milko.videv@thalesgroup.com)
-----
Last Modified: Tuesday, 5th March 2024 7:00:29 pm
Modified By: Milko Videv (milko.videv@thalesgroup.com>)
-----
Copyright 2017 - 2024, Thales DIS, MCS SSH
-----
HISTORY:
Date      	By	Comments
----------	---	---------------------------------------------------------
'''

from fastcore.all import *
from fastai.vision.all import *
from tools import *
import pathlib

plt = platform.system()
if plt != 'Windows': pathlib.WindowsPath = pathlib.PosixPath

learn = load_learner("./models/trained.pkl")
labels = learn.dls.vocab

def predict(image):    
    if image is None:
        return None
    img = PILImage.create(image)
    pred,pred_idx,probs = learn.predict(img)
    result = {labels[i]: float(probs[i]) for i in range(len(labels))}

    labels_map = {'bulgarian': 'българска', 'indian': 'индийска', 'japanese': 'японска'}
    result = {labels_map[key]: value for key, value in result.items()}

    max_type = max(result, key=result.get)
    max_probability = result[max_type] * 100    
    print(f"Тази шевица е {max_type} с вероятност {max_probability:.2f}%")

    return result

if __name__ == "__main__":
    if debugger_is_active():
        predict("./samples/bulgarian.jpg")
        sys.exit()
    else:
        if len(sys.argv) < 2:
            print("Use: python predict.py <image path>")   
            print("Example: python predict.py samples/bulgarian.jpg")
            sys.exit()

        src = sys.argv[1]    
        predict(src)