Robustness Update: Added confidence thresholding, file type validation, and error handling

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 import torch
 import torch.nn as nn
 from torchvision import models, transforms
@@ -16,9 +16,12 @@ num_ftrs = model.head.in_features
 model.head = nn.Linear(num_ftrs, 2)
 
 # Load Swin weights
-model.load_state_dict(torch.load("models/spoilage_model_swin_finetuned.pth", map_location=device))
-model = model.to(device)
-model.eval()
+try:
+    model.load_state_dict(torch.load("models/spoilage_model_swin_finetuned.pth", map_location=device))
+    model = model.to(device)
+    model.eval()
+except Exception as e:
+    print(f"Error loading model: {e}")
 
 # Preprocessing transforms
 preprocess = transforms.Compose([
@@ -29,35 +32,59 @@ preprocess = transforms.Compose([
 ])
 
 CLASS_NAMES = ["fresh", "spoiled"]
+ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg"}
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # Read image
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    
-    # Preprocess
-    input_tensor = preprocess(image)
-    input_batch = input_tensor.unsqueeze(0).to(device)
-    
-    # Inference
-    with torch.no_grad():
-        outputs = model(input_batch)
-        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-        confidence, index = torch.max(probabilities, 0)
-    
-    result = {
-        "prediction": CLASS_NAMES[index.item()],
-        "confidence": float(confidence.item()),
-        "status": "success",
-        "model": "swin_transformer_tiny"
-    }
-    
-    return result
+    # 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}")
+
+    try:
+        # 2. Read and Validate Image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        
+        # 3. Preprocess
+        input_tensor = preprocess(image)
+        input_batch = input_tensor.unsqueeze(0).to(device)
+        
+        # 4. Inference
+        with torch.no_grad():
+            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",
+                "message": "Model confidence too low. The image might not be a recognized food item or is too low quality.",
+                "confidence": conf_score,
+                "suggestion": "Please provide a clearer image of a fruit or vegetable."
+            }
+        
+        return {
+            "prediction": CLASS_NAMES[index.item()],
+            "confidence": conf_score,
+            "status": "success",
+            "model": "swin_transformer_tiny_finetuned"
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
 @app.get("/")
 def read_root():
-    return {"message": "Food Spoilage Detection API (Swin) is running. Use /predict to send an image."}
+    return {
+        "message": "Food Spoilage Detection API is running.",
+        "usage": "POST an image to /predict",
+        "supported_classes": CLASS_NAMES
+    }
 
 if __name__ == "__main__":
     import uvicorn

deploy/app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 import torch
 import torch.nn as nn
 from torchvision import models, transforms
@@ -16,9 +16,12 @@ num_ftrs = model.head.in_features
 model.head = nn.Linear(num_ftrs, 2)
 
 # Load Swin weights
-model.load_state_dict(torch.load("models/spoilage_model_swin_finetuned.pth", map_location=device))
-model = model.to(device)
-model.eval()
+try:
+    model.load_state_dict(torch.load("models/spoilage_model_swin_finetuned.pth", map_location=device))
+    model = model.to(device)
+    model.eval()
+except Exception as e:
+    print(f"Error loading model: {e}")
 
 # Preprocessing transforms
 preprocess = transforms.Compose([
@@ -29,35 +32,59 @@ preprocess = transforms.Compose([
 ])
 
 CLASS_NAMES = ["fresh", "spoiled"]
+ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg"}
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # Read image
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    
-    # Preprocess
-    input_tensor = preprocess(image)
-    input_batch = input_tensor.unsqueeze(0).to(device)
-    
-    # Inference
-    with torch.no_grad():
-        outputs = model(input_batch)
-        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-        confidence, index = torch.max(probabilities, 0)
-    
-    result = {
-        "prediction": CLASS_NAMES[index.item()],
-        "confidence": float(confidence.item()),
-        "status": "success",
-        "model": "swin_transformer_tiny"
-    }
-    
-    return result
+    # 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}")
+
+    try:
+        # 2. Read and Validate Image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        
+        # 3. Preprocess
+        input_tensor = preprocess(image)
+        input_batch = input_tensor.unsqueeze(0).to(device)
+        
+        # 4. Inference
+        with torch.no_grad():
+            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",
+                "message": "Model confidence too low. The image might not be a recognized food item or is too low quality.",
+                "confidence": conf_score,
+                "suggestion": "Please provide a clearer image of a fruit or vegetable."
+            }
+        
+        return {
+            "prediction": CLASS_NAMES[index.item()],
+            "confidence": conf_score,
+            "status": "success",
+            "model": "swin_transformer_tiny_finetuned"
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
 @app.get("/")
 def read_root():
-    return {"message": "Food Spoilage Detection API (Swin) is running. Use /predict to send an image."}
+    return {
+        "message": "Food Spoilage Detection API is running.",
+        "usage": "POST an image to /predict",
+        "supported_classes": CLASS_NAMES
+    }
 
 if __name__ == "__main__":
     import uvicorn