Update newapi.py

newapi.py

@@ -1,51 +1,100 @@
-from fastapi import FastAPI, UploadFile, File
+import os
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
-from utils import BrainTumorModel, get_precautions_from_gemini
 
-app = FastAPI()
+from utils import BrainTumorModel, GliomaStageModel, get_precautions_from_gemini
 
-# Load the model
-btd_model = BrainTumorModel()
-btd_model_path = "brain_tumor_model.pth"
+# ---- Constants ----
+MODEL_DIR = "models"
+BTD_FILENAME   = "BTD_model.pth"
+GLIO_FILENAME = "glioma_stages.pth"
 
-try:
-    btd_model.load_state_dict(torch.load(btd_model_path, map_location=torch.device('cpu')))
-    btd_model.eval()
-except Exception as e:
-    print(f"❌ Error loading model: {e}")
+# ---- App setup ----
+app = FastAPI(title="Brain Tumor Detection API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # adjust in production
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# Define image transform
+# ---- Device & transforms ----
+DEVICE = torch.device("cpu")
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
 ])
 
-# Class labels (adjust if your model uses different labels)
-classes = ['glioma', 'meningioma', 'notumor', 'pituitary']
+# ---- Load & init models ----
+def load_model(cls, filename):
+    path = os.path.join(MODEL_DIR, filename)
+    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
+
+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 ----
 
 @app.get("/")
-def read_root():
-    return {"message": "Brain Tumor Detection API is running 🚀"}
-
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    try:
-        image = Image.open(file.file).convert("RGB")
-        image = transform(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]
-
-        with torch.no_grad():
-            outputs = btd_model(image)
-            _, predicted = torch.max(outputs.data, 1)
-            predicted_class = classes[predicted.item()]
-            precautions = get_precautions_from_gemini(predicted_class)
-
-        return JSONResponse(content={
-            "prediction": predicted_class,
-            "precautions": precautions
-        })
-
-    except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=500)
+async def health():
+    return {"status": "ok", "message": "Brain Tumor API is live"}
+
+@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")
+    img = Image.open(file.file).convert("RGB")
+    tensor = 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']
+    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)
+        }
+
+class MutationInput(BaseModel):
+    gender: str
+    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)
+
+    with torch.no_grad():
+        out = glioma_model(tensor)
+        idx = torch.argmax(out, dim=1).item()
+
+    stages = ['Stage 1','Stage 2','Stage 3','Stage 4']
+    return {"glioma_stage": stages[idx]}