image processing and prediction

app.py

@@ -1,15 +1,30 @@
-# import keras
+from keras import (
+    saving,
+    preprocessing,
+    applications
+)
 import fastapi
 from fastapi import UploadFile, File, HTTPException
 from PIL import Image
 import io
 import time
+import numpy as np
+
 
 app = fastapi.FastAPI()
 
-# model = keras.saving.load_model("hf://Yumeng-Liu/trash-classifier")
+model = saving.load_model("hf://Yumeng-Liu/trash-classifier")
+
 
-print("Started")
+def get_prediction(img: Image):
+    img = img.resize((224, 224))
+    img = img.convert("L")
+    img_array = preprocessing.image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Add an extra dimension to match the model's input shape
+    img_array = applications.mobilenet_v2.preprocess_input(img_array)
+
+    prediction = model.predict(img_array)
+    return prediction
 
 
 @app.get("/")
@@ -24,11 +39,18 @@ async def predict(received_image: UploadFile = File(...)):
         # Open the binary data as an image
         image = Image.open(io.BytesIO(contents))
 
+        print("Image received")
         # You can now work with the `image` object
         print(image.format, image.size, image.mode)  # Example: JPEG (1920, 1080) RGB
+        print("")
+
+        prediction_result = get_prediction(image)
+        print(prediction_result)
 
         # Perform further processing, e.g., save it, analyze it, etc.
-        return {"filename": received_image.filename, "message": "Image processed successfully"}
+        return {
+            "result": prediction_result
+        }
 
     except Exception as e:
         print(e)
@@ -39,5 +61,6 @@ async def predict(received_image: UploadFile = File(...)):
 
 
 if __name__ == "__main__":
+    print("Starting app")
     while True:
         time.sleep(10)

requirements.txt

@@ -4,4 +4,4 @@ tensorflow
 keras
 huggingface_hub
 python-multipart
-pillow
+pillow