Update app.py

app.py

@@ -4,6 +4,11 @@ import keras
 import h5py
 import numpy as np
 
+from pydantic import BaseModel
+
+class Prediction(BaseModel):
+    prediction: np.ndarray
+
 app = FastAPI()
 
 # Load the model from the HDF5 file
@@ -11,24 +16,17 @@ model = keras.models.load_model("model_finetuned.h5")
 
 # Define a function to preprocess the image
 def preprocess_image(image):
-    # Resize the image to a fixed size
-    image = image.resize((224, 224))
-    # Convert the image to a NumPy array
-    image = np.array(image)
-    # Normalize the image
-    image = image / 255.0
-    # Return the image
-    return image
+    # ...
 
 # Define an endpoint to predict the output
-@app.post("/predict")
+@app.post("/predict", response_model=Prediction)
 async def predict_endpoint(image: fastapi.File):
     # Preprocess the image
     image = preprocess_image(image)
     # Make a prediction
     prediction = model.predict(image)
     # Return the prediction
-    return {"prediction": prediction}
+    return Prediction(prediction=prediction)
 
 if __name__ == "__main__":
     import uvicorn