fix: add file size limit, return 400 on invalid images, use inference_mode

app.py

@@ -2,8 +2,10 @@ from fastapi import FastAPI, File, UploadFile, HTTPException
 import torch
 import torch.nn as nn
 from torchvision import models, transforms
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import io
+import uvicorn
+
 
 app = FastAPI(title="Food Spoilage Detection API (Swin)")
 
@@ -29,34 +31,31 @@ preprocess = transforms.Compose([
 ])
 
 CLASS_NAMES = ["fresh", "spoiled"]
-ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg"}
+MAX_FILE_SIZE = 20 * 1024 * 1024  # 10 MB
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # 1. Validate File Extension
-    extension = file.filename.split(".")[-1].lower()
-    if extension not in ALLOWED_EXTENSIONS:
-        raise HTTPException(status_code=400, detail=f"Invalid file type. Allowed: {ALLOWED_EXTENSIONS}")
+    contents = await file.read()
+
+    if len(contents) > MAX_FILE_SIZE:
+        raise HTTPException(status_code=400, detail="File too large. Maximum size is 10 MB.")
 
     try:
-        # 2. Read and Validate Image
-        contents = await file.read()
         image = Image.open(io.BytesIO(contents)).convert("RGB")
-        
-        # 3. Preprocess
+    except UnidentifiedImageError:
+        raise HTTPException(status_code=400, detail="Could not read image. Please upload a valid image file.")
+
+    try:
         input_tensor = preprocess(image)
         input_batch = input_tensor.unsqueeze(0).to(device)
-        
-        # 4. Inference
-        with torch.no_grad():
+
+        with torch.inference_mode():
             outputs = model(input_batch)
             probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
             confidence, index = torch.max(probabilities, 0)
-        
+
         conf_score = float(confidence.item())
-        
-        # 5. Robustness Check: Confidence Threshold
-        # If the model is not confident enough, it might not be a food image
+
         if conf_score < 0.70:
             return {
                 "status": "ambiguous",
@@ -64,7 +63,7 @@ async def predict(file: UploadFile = File(...)):
                 "confidence": conf_score,
                 "suggestion": "Please provide a clearer image of a fruit or vegetable."
             }
-        
+
         return {
             "prediction": CLASS_NAMES[index.item()],
             "confidence": conf_score,
@@ -84,5 +83,4 @@ def read_root():
     }
 
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)