Create app.py

app.py

@@ -0,0 +1,38 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+import torch
+from transformers import AutoTokenizer, AutoModelForImageClassification
+from PIL import Image
+import torchvision.transforms as T
+import requests
+from io import BytesIO
+
+app = FastAPI()
+
+# load model once
+model_name = "nateraw/vit-base-patch16-224-in21k"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+transform = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize([0.5], [0.5])
+])
+
+class ImageInput(BaseModel):
+    url: str
+
+@app.get("/")
+def read_root():
+    return {"status": "running"}
+
+@app.post("/predict")
+def predict(input: ImageInput):
+    img = Image.open(BytesIO(requests.get(input.url).content)).convert("RGB")
+    img_tensor = transform(img).unsqueeze(0)
+
+    with torch.no_grad():
+        logits = model(img_tensor).logits
+        pred = torch.argmax(logits, dim=1).item()
+
+    return {"class": int(pred)}