Update app.py

app.py

@@ -4,15 +4,13 @@ from fastapi.templating import Jinja2Templates
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 import torch
-from torchvision import models, transforms
+from torchvision import transforms, models
 from PIL import Image
 import io
-import torch.nn as nn
 
-# FastAPI app setup
 app = FastAPI()
 
-# CORS middleware (optional for frontend access)
+# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -28,33 +26,37 @@ templates = Jinja2Templates(directory="templates")
 # Device
 device = torch.device("cpu")
 
-# Model architecture
+# Model (assumes ResNet18 was used)
 model = models.resnet18(weights=None)
-model.fc = nn.Linear(model.fc.in_features, 2)  # Binary classification
+model.fc = torch.nn.Linear(model.fc.in_features, 2)  # Assuming 2 classes: Normal & Pneumonia
 model.load_state_dict(torch.load("pneumonia_weights.pth", map_location=device))
+model.to(device)
 model.eval()
 
-# Image transform
+# Transform
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
-# Homepage route
+# Routes
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
-# Prediction route
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    image = transform(image).unsqueeze(0).to(device)
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        image = transform(image).unsqueeze(0).to(device)
 
-    with torch.no_grad():
-        output = model(image)
-        prediction = torch.argmax(output, 1).item()
-        result = "Pneumonia" if prediction == 1 else "Normal"
+        with torch.no_grad():
+            output = model(image)
+            predicted = torch.argmax(output, dim=1).item()
+            result = "Pneumonia" if predicted == 1 else "Normal"
 
-    return {"result": result}
+        return {"result": result}
+
+    except Exception as e:
+        return {"result": f"Error during prediction: {str(e)}"}