Update app.py

app.py

@@ -2,22 +2,16 @@ import fastapi
 from transformers import pipeline
 import h5py
 
-# Load the model from the h5 file
-model = h5py.File("model_finetuned.h5", "r")["model"]
+# Create a context manager to load the model from the HDF5 file
+with h5py.File("model_finetuned.h5", "r") as f:
+    model = f["model"]
 
-# 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 a dependency to inject the model into the `predict_endpoint()` function
+def get_model():
+    return model
 
 # Define an endpoint to predict the output
-@app.post("/predict")
+@app.post("/predict", dependencies=[Depends(get_model)])
 async def predict_endpoint(image: fastapi.File):
     # Preprocess the image
     image = preprocess_image(image)
@@ -28,6 +22,14 @@ async def predict_endpoint(image: fastapi.File):
     # Return the prediction
     return {"prediction": prediction}
 
+# Define a FastAPI exception handler to handle any errors that occur while making predictions
+@app.exception_handler(Exception)
+async def handle_exception(request: Request, exc: Exception):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"message": str(exc)},
+    )
+
 # Start the FastAPI app
 if __name__ == "__main__":
     import uvicorn