Spaces:
Running
Running
File size: 1,763 Bytes
30c5114 df7fc7d 30c5114 df7fc7d 30c5114 6fc2183 df7fc7d 30c5114 df7fc7d 30c5114 df7fc7d 30c5114 df7fc7d 30c5114 df7fc7d 30c5114 ed89f9e 30c5114 ed89f9e 30c5114 df7fc7d 30c5114 ed89f9e 30c5114 | 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 49 50 51 52 53 54 55 56 57 58 59 60 | from utils.preprocessing import preprocess_image
import tensorflow as tf
import numpy as np
import os
import logging
import json
class AfroPalmModel:
"""
Class to load the model and make predictions
"""
def __init__(self):
logging.info("Loading classification model")
self.model_path = os.path.dirname(os.path.abspath("ghostnet_model_float32.tflite")) + "/models/ghostnet_model_float32.tflite"
logging.debug(f"Preparing to read from {self.model_path}")
with open(self.model_path, 'rb') as fid:
tflite_model = fid.read()
logging.info("File read successfully")
# Create and allocate the interpreter using the loaded state
self.interpreter = tf.lite.Interpreter(model_content=tflite_model)
self.interpreter.allocate_tensors()
# Retrieve the input and output indices
self.input_index = self.interpreter.get_input_details()[0]["index"]
self.output_index = self.interpreter.get_output_details()[0]["index"]
logging.info("Model loaded successfully")
def predict(self, image_path):
"""
Make a prediction on the image
:param image: image to make prediction on
:return: prediction and confidence score
"""
logging.info("Making prediction")
img = preprocess_image(image_path)
logging.debug(f'Image preprocessed with shape {np.array(img).shape}')
self.interpreter.set_tensor(self.input_index, img)
self.interpreter.invoke()
predictions = list(self.interpreter.get_tensor(self.output_index)[0])
logging.info("Classification successful")
return predictions.index(max(predictions)), max(predictions) |