api/predictor.py

import json
import os

import numpy as np
import tensorflow as tf

from api.preprocess import preprocess_image

from config import (
    SAVED_MODELS_DIR,
    MODEL_KERAS_NAME,
    CLASS_INDEX_FILE,
)

MODEL_PATH = os.path.join(
    SAVED_MODELS_DIR,
    MODEL_KERAS_NAME,
)

print("Loading TensorFlow model...")

model = tf.keras.models.load_model(MODEL_PATH)

print("Model Loaded")

with open(CLASS_INDEX_FILE) as f:
    class_indices = json.load(f)

index_to_class = {
    v: k
    for k, v in class_indices.items()
}


def predict(image):

    image = preprocess_image(image)

    prediction = model.predict(image, verbose=0)

    class_id = np.argmax(prediction)

    confidence = float(np.max(prediction))

    return {
        "label": index_to_class[class_id],
        "confidence": round(confidence * 100, 2),
    }