/ src / model / model.py
model.py
 1  import io
 2  import numpy as np
 3  import os
 4  from PIL import Image
 5  import tensorflow as tf
 6  
 7  MODEL_NAME = "model.keras"
 8  
 9  
10  def preprocess_image_bytes(bytes):
11      image = Image.open(io.BytesIO(bytes))
12      image = np.array(image)
13      image = np.expand_dims(image, axis=0)
14  
15      return image
16  
17  
18  def process_model_bytes(bytes):
19      open(MODEL_NAME, "wb").write(bytes)
20      model = tf.keras.models.load_model(MODEL_NAME)
21      os.remove(MODEL_NAME)
22  
23      return model
24  
25  
26  def predict(model_bytes, image_bytes):
27      image = preprocess_image_bytes(image_bytes)
28      model = process_model_bytes(model_bytes)
29      result = model.predict(image)
30  
31      return result[0][0], result[0][1]