Create app.py

app.py

@@ -0,0 +1,76 @@
+import io
+from typing import List
+from fastapi.middleware.cors import CORSMiddleware
+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)")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all domains. For security, replace with your GitHub Pages URL later.
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+async def root():
+    return {"message": "Image Captioning API is running. Use /docs for the UI or POST /caption for captions."}
+
+@app.post("/caption")
+async def get_caption(file: UploadFile = File(...)):
+    # Your existing logic to process the image and generate a caption
+    # result = model.predict(image)
+    return {"caption": "The generated caption text here"}
+    
+device = get_device()
+training_cfg = TrainingConfig(max_caption_length=50)
+tokenizer = create_tokenizer()
+model = ImageCaptioningModel(training_cfg=training_cfg)
+model.to(device)
+model.eval()
+
+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,
+        )
+
+    return JSONResponse({"caption": captions[0]})