Update app.py

app.py

@@ -9,10 +9,10 @@ from typing import Any, Type
 import torch
 
 class TorchTensor(torch.Tensor):
-  pass
+ pass
 
-class Prediction(BaseModel):
-  prediction: TorchTensor
+class Prediction():
+ prediction: TorchTensor
 
 app = FastAPI()
 
@@ -23,55 +23,50 @@ model = torch.load("best_model-epoch=01-val_loss=3.00.ckpt")
 
 
 def preprocess_input(input):
-  """Preprocess the input image for the PyTorch image classification model.
+ """Preprocess the input image for the PyTorch image classification model.
+ Args:
+  input: A PIL Image object.
+ Returns:
+  A PyTorch tensor representing the preprocessed image.
+ """
 
-  Args:
-    input: A PIL Image object.
+ # Resize the image to the expected size.
+ input = input.resize((224, 224))
 
-  Returns:
-    A PyTorch tensor representing the preprocessed image.
-  """
+ # Convert the image to a PyTorch tensor.
+ input = torch.from_numpy(np.array(input)).float()
 
-  # Resize the image to the expected size.
-  input = input.resize((224, 224))
+ # Normalize the image.
+ input = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input)
 
-  # Convert the image to a PyTorch tensor.
-  input = torch.from_numpy(np.array(input)).float()
+ # Return the preprocessed image.
+ return input
 
-  # Normalize the image.
-  input = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input)
+@app.post("/predict")
+async def predict_endpoint(input: Any):
+ """Predict the output of the PyTorch image classification model.
+ Args:
+  input: A file containing the input image.
+ Returns:
+  A JSON object containing the prediction.
+ """
 
-  # Return the preprocessed image.
-  return input
+ # Load the image.
+ image = Image.open(BytesIO(input))
 
-@app.post("/predict", response_model=Prediction)
-async def predict_endpoint(input: fastapi.File):
-  """Predict the output of the PyTorch image classification model.
+ # Preprocess the image.
+ image = preprocess_input(image)
 
-  Args:
-    input: A file containing the input image.
+ # Make a prediction.
+ prediction = model(image.unsqueeze(0))
 
-  Returns:
-    A JSON object containing the prediction.
-  """
+ # Get the top predicted class.
+ predicted_class = prediction.argmax(1)
 
-  # Load the image.
-  image = await input.read()
-  image = Image.open(BytesIO(image))
+ # Return the prediction.
+ return {"prediction": predicted_class.item()}
 
-  # Preprocess the image.
-  image = preprocess_input(image)
 
-  # Make a prediction.
-  prediction = model(image.unsqueeze(0))
-
-  # Get the top predicted class.
-  predicted_class = prediction.argmax(1)
-
-  # Return the prediction.
-  return {"prediction": predicted_class.item()}
-
-
-if __name__ == "__main__":
-  import uvicorn
-  uvicorn.run(app, host="0.0.0.0", port=8000)
+if _name_ == "_main_":
+ import uvicorn
+ uvicorn.run(app, host="0.0.0.0", port=8000)