Update newapi.py

newapi.py

@@ -2,99 +2,89 @@ import os
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
 
 from utils import BrainTumorModel, GliomaStageModel, get_precautions_from_gemini
 
-# ---- Constants ----
-MODEL_DIR = "models"
-BTD_FILENAME   = "BTD_model.pth"
+# ——— Constants ———
+MODEL_DIR     = "models"
+BTD_FILENAME  = "BTD_model.pth"
 GLIO_FILENAME = "glioma_stages.pth"
 
-# ---- App setup ----
+# ——— App setup ———
 app = FastAPI(title="Brain Tumor Detection API")
-
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # adjust in production
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# ---- Device & transforms ----
 DEVICE = torch.device("cpu")
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
+    transforms.Normalize([0.5]*3, [0.5]*3),
 ])
 
-# ---- Load & init models ----
-def load_model(cls, filename):
-    path = os.path.join(MODEL_DIR, filename)
+def load_model(cls, fname):
+    path = os.path.join(MODEL_DIR, fname)
     if not os.path.isfile(path):
         raise FileNotFoundError(f"Model file not found: {path}")
-    model = cls().to(DEVICE).eval()
-    model.load_state_dict(torch.load(path, map_location=DEVICE))
-    return model
+    m = cls().to(DEVICE)
+    m.load_state_dict(torch.load(path, map_location=DEVICE))
+    return m.eval()
 
 try:
     tumor_model  = load_model(BrainTumorModel, BTD_FILENAME)
     glioma_model = load_model(GliomaStageModel, GLIO_FILENAME)
 except Exception as e:
-    # During startup, any exception here will show in logs
-    print(f"❌ Error loading model: {e}")
-
-# ---- Routes ----
+    print("❌ Error loading model:", e)
 
 @app.get("/")
 async def health():
-    return {"status": "ok", "message": "Brain Tumor API is live"}
+    return {"status": "ok", "message": "API is up"}
 
 @app.post("/predict-image/")
 async def predict_image(file: UploadFile = File(...)):
-    if file.content_type.split("/")[0] != "image":
-        raise HTTPException(400, "Upload an image file")
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(400, "Upload an image")
     img = Image.open(file.file).convert("RGB")
-    tensor = transform(img).unsqueeze(0).to(DEVICE)
+    t   = transform(img).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        out = tumor_model(tensor)
-        idx = torch.argmax(out, dim=1).item()
-
-    labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
+        out = tumor_model(t)
+        idx = int(out.argmax(1))
+    labels = ["glioma","meningioma","notumor","pituitary"]
     tumor_type = labels[idx]
-
     if tumor_type == "glioma":
         return {"tumor_type": tumor_type, "next": "submit_mutation_data"}
-    else:
-        return {
-            "tumor_type": tumor_type,
-            "precaution": get_precautions_from_gemini(tumor_type)
-        }
+    return {
+        "tumor_type": tumor_type,
+        "precaution": get_precautions_from_gemini(tumor_type)
+    }
 
 class MutationInput(BaseModel):
     gender: str
-    age: float
-    idh1: int
-    tp53: int
-    atrx: int
-    pten: int
-    egfr: int
-    cic: int
+    age:    float
+    idh1:   int
+    tp53:   int
+    atrx:   int
+    pten:   int
+    egfr:   int
+    cic:    int
     pik3ca: int
 
 @app.post("/predict-glioma-stage/")
 async def predict_glioma_stage(data: MutationInput):
-    gender_val = 0 if data.gender.lower().startswith('m') else 1
-    features = [gender_val, data.age, data.idh1, data.tp53,
-                data.atrx, data.pten, data.egfr, data.cic, data.pik3ca]
-    tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(DEVICE)
-
+    gen = 0 if data.gender.lower().startswith("m") else 1
+    feats = [gen, data.age, data.idh1, data.tp53,
+             data.atrx, data.pten, data.egfr, data.cic, data.pik3ca]
+    t = torch.tensor(feats, dtype=torch.float32).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        out = glioma_model(tensor)
-        idx = torch.argmax(out, dim=1).item()
-
-    stages = ['Stage 1','Stage 2','Stage 3','Stage 4']
+        out = glioma_model(t)
+        idx = int(out.argmax(1))
+    stages = ["Stage 1","Stage 2","Stage 3","Stage 4"]
     return {"glioma_stage": stages[idx]}