Update inference.py

inference.py

@@ -1,24 +1,32 @@
 from PIL import Image
 import tensorflow as tf
 import numpy as np
+import io
+import base64
 
-# Load the TFLite model
+# Load the TFLite model once
 interpreter = tf.lite.Interpreter(model_path="leaf_model_85_percent.tflite")
 interpreter.allocate_tensors()
 input_details = interpreter.get_input_details()
 output_details = interpreter.get_output_details()
 
-# Define preprocess and inference
 def preprocess(image: Image.Image):
     img = image.resize((224, 224)).convert("RGB")
     img_array = np.array(img) / 255.0
     img_array = np.expand_dims(img_array, axis=0).astype(np.float32)
     return img_array
 
-def predict(image: Image.Image):
+def classify(image_bytes: bytes):
+    image = Image.open(io.BytesIO(image_bytes))
     input_tensor = preprocess(image)
     interpreter.set_tensor(input_details[0]['index'], input_tensor)
     interpreter.invoke()
     output = interpreter.get_tensor(output_details[0]['index'])[0][0]
     label = "Unhealthy" if output > 0.5 else "Healthy"
-    return [{"label": label, "score": float(output)}]
+    return {"label": label, "score": float(output)}
+
+# This is the function HF Spaces calls to do inference
+def inference(payload):
+    # payload will have image data in base64 or bytes format, depends on your API input
+    image_bytes = base64.b64decode(payload["data"][0].split(",")[1])  # Assuming base64 input
+    return classify(image_bytes)