File size: 824 Bytes
f0e10d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | 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),
} |