Update app.py

app.py

@@ -1,39 +1,46 @@
+import torch
 import fastapi
-from transformers import pipeline
-import pickle
+import numpy as np
+from PIL import Image
+from typing import Any, Type
 
-# Load the model from the pickle file
-with open("model2.ckpt", "rb") as f:
-    model = pickle.load(f)
+class TorchTensor(torch.Tensor):
+    pass
 
-# Define a function to preprocess the image
-def preprocess_image(image):
-    # Resize the image to a fixed size
-    image = image.resize((224, 224))
+class Prediction:
+    prediction: TorchTensor
 
-    # Convert the image to a NumPy array
-    image = np.array(image)
+app = fastapi.FastAPI()
 
-    # Normalize the image
-    image = image / 255.0
+model = torch.load("model67.bin", map_location='cpu')
 
-    # Return the image
-    return image
+# Define a function to preprocess the input image
+def preprocess_input(input: Any):
+    image = Image.open(BytesIO(input))
+    image = image.resize((224, 224))
+    input = np.array(image)
+    input = torch.from_numpy(input).float()
+    input = input.permute(2, 0, 1)
+    input = input.unsqueeze(0)
+    return input
 
-# Define an endpoint to predict the output
+# Define an endpoint to make predictions
 @app.post("/predict")
-async def predict_endpoint(image: fastapi.File):
-    # Preprocess the image
-    image = preprocess_image(image)
+async def predict_endpoint(input: Any):
+    """Make a prediction on an image uploaded by the user."""
+
+    # Preprocess the input image
+    input = preprocess_input(input)
 
     # Make a prediction
-    prediction = model(image)
+    prediction = model(input)
 
-    # Return the prediction
-    return {"prediction": prediction}
+    # Get the predicted class
+    predicted_class = prediction.argmax(1).item()
 
-# Start the FastAPI app
-if _name_ == "_main_":
-    import uvicorn
+    # Return the predicted class in JSON format
+    return {"prediction": predicted_class}
 
+if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)