Spaces:
Runtime error
Runtime error
# app.py | |
from fastapi import FastAPI | |
import torch | |
import torch.nn as nn | |
import torch | |
from torchvision import transforms | |
from typing import Any, Type | |
import torch | |
class TorchTensor(torch.Tensor): | |
pass | |
class Prediction: | |
prediction: TorchTensor | |
app = FastAPI() | |
# Load the PyTorch model | |
model = torch.load("best_model-epoch=01-val_loss=3.00.ckpt") | |
# Define a function to preprocess the input | |
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) | |
return input | |
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()} | |
if _name_ == "_main_": | |
import uvicorn | |
uvicorn.run(app, host="0.0.0.0", port=8000) | |