Update app.py

app.py

@@ -25,50 +25,31 @@ 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.
- Args:
- input: A PIL Image object.
- Returns:
- A PyTorch tensor representing the preprocessed image.
- """
-
- # Resize the image to the expected size.
- input = input.resize((224, 224))
-
- # Convert the image to a PyTorch tensor.
- input = torch.from_numpy(np.array(input)).float()
-
- # Normalize the image.
- input = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input)
-
- # Return the preprocessed image.
- return 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)
+    
+    return 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.
- """
-
- # Load the image.
- image = Image.open(BytesIO(input))
-
- # 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()}
+    
+    
+    image = Image.open(BytesIO(input))
+    
+    image = preprocess_input(image)
+    
+    prediction = model(image.unsqueeze(0))
+    
+    predicted_class = prediction.argmax(1)
+    
+    return {"prediction": predicted_class.item()}
 
 
 if _name_ == "_main_":
- import uvicorn
- uvicorn.run(app, host="0.0.0.0", port=8000)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)