Update app.py

app.py

@@ -1,51 +1,47 @@
-# app.py
-
-from fastapi import FastAPI
-import torch
-import torch.nn as nn
 import torch
-from torchvision import transforms
-
+import fastapi
+import numpy as np
+from PIL import Image
 from typing import Any, Type
 
-import torch
-
 class TorchTensor(torch.Tensor):
     pass
 
 class Prediction:
     prediction: TorchTensor
 
-app = FastAPI()
+app = fastapi.FastAPI()
 
-model = torch.load("/main/best_model-epoch=01-val_loss=3.00.ckpt")
+# Load the .bin model
+model = torch.load("/main/best_model-epoch=01-val_loss=3.00.ckpt", map_location='cpu')
 
-def preprocess_input(input):
-    
-    
-    input = input.resize((224, 224))
-    
-    input = torch.from_numpy(np.array(input)).float()
-    
-    input = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input)
-    
+# 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 make predictions
 @app.post("/predict")
 async def predict_endpoint(input: Any):
-    
-    
-    image = Image.open(BytesIO(input))
-    
-    image = preprocess_input(image)
-    
-    prediction = model(image.unsqueeze(0))
-    
-    predicted_class = prediction.argmax(1)
-    
-    return {"prediction": predicted_class.item()}
+    """Make a prediction on an image uploaded by the user."""
+
+    # Preprocess the input image
+    input = preprocess_input(input)
+
+    # Make a prediction
+    prediction = model(input)
+
+    # Get the predicted class
+    predicted_class = prediction.argmax(1).item()
 
+    # Return the predicted class in JSON format
+    return {"prediction": predicted_class}
 
-if _name_ == "_main_":
+if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)