Update app.py

app.py

@@ -11,66 +11,65 @@ 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
 
+# 1. Initialize App and CORS
 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_origins=["*"],
     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"}
-    
+# 2. Load Model & Assets (Global Scope)
 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 weights
 CHECKPOINT_PATH = "best_model.pt"
 state_dict = torch.load(CHECKPOINT_PATH, map_location=device)
 model.load_state_dict(state_dict)
+model.to(device)
+model.eval()
 
-preprocess = transforms.Compose(
-    [
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
-    ]
-)
+# 3. Preprocessing Pipeline
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+])
 
+# 4. API Routes
+@app.get("/")
+async def root():
+    return {"message": "API is online. Go to /docs for testing."}
 
 @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}"})
+        
+        # Preprocess and Move to Device
+        tensor = preprocess(image).unsqueeze(0).to(device)
 
-    tensor = preprocess(image).unsqueeze(0).to(device)
+        # Inference
+        with torch.no_grad():
+            captions: List[str] = model.generate(
+                images=tensor,
+                max_length=50,
+                num_beams=1,
+            )
 
-    with torch.no_grad():
-        captions: List[str] = model.generate(
-            images=tensor,
-            max_length=50,
-            num_beams=1,
-        )
+        return JSONResponse({"caption": captions[0]})
 
-    return JSONResponse({"caption": captions[0]})
+    except Exception as exc:
+        return JSONResponse(status_code=400, content={"error": f"Internal Error: {exc}"})