Create app.py

image_captioning/app.py

@@ -0,0 +1,60 @@
+import io
+from typing import List
+
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from PIL import Image
+import torch
+from torchvision import transforms
+
+from image_captioning.config import TrainingConfig, get_device
+from image_captioning.dataset import IMAGENET_MEAN, IMAGENET_STD, create_tokenizer
+from image_captioning.model import ImageCaptioningModel
+
+app = FastAPI(title="Image Captioning API (HF Space)")
+
+device = get_device()
+training_cfg = TrainingConfig(max_caption_length=50)
+tokenizer = create_tokenizer()
+model = ImageCaptioningModel(training_cfg=training_cfg)
+model.to(device)
+model.eval()
+
+# Load checkpoint from the repo root
+CHECKPOINT_PATH = "best_model.pt"
+state_dict = torch.load(CHECKPOINT_PATH, map_location=device)
+model.load_state_dict(state_dict)
+
+preprocess = transforms.Compose(
+    [
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+    ]
+)
+
+
+@app.get("/health")
+async def health() -> dict:
+    return {"status": "ok"}
+
+
+@app.post("/caption")
+async def caption_image(file: UploadFile = File(...)) -> JSONResponse:
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+    except Exception as exc:
+        return JSONResponse(status_code=400, content={"error": f"Invalid image: {exc}"})
+
+    tensor = preprocess(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        captions: List[str] = model.generate(
+            images=tensor,
+            max_length=50,
+            num_beams=1,  # deterministic greedy decoding
+        )
+
+    return JSONResponse({"caption": captions[0]})