Update newapi.py

newapi.py

@@ -1,21 +1,22 @@
-import os
-from fastapi import FastAPI, File, UploadFile, HTTPException
+# newapi.py
+
+from fastapi import FastAPI, File, UploadFile
 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 PIL import Image
+import io
+import os
 
-from utils import BrainTumorModel, GliomaStageModel, get_precautions_from_gemini
+# Use a writable directory in Hugging Face Spaces
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache"
+os.environ["HF_HOME"] = "/tmp/.cache"
 
-# ——— Constants ———
-MODEL_DIR     = "models"
-BTD_FILENAME  = "BTD_model.pth"
-GLIO_FILENAME = "glioma_stages.pth"
+# Define FastAPI app
+app = FastAPI(title="🧠 Brain Tumor Detection API")
 
-# ——— App setup ———
-app = FastAPI(title="Brain Tumor Detection API")
+# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -23,68 +24,62 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-DEVICE = torch.device("cpu")
+# Image transform
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5]*3, [0.5]*3),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
 ])
 
-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}")
-    m = cls().to(DEVICE)
-    m.load_state_dict(torch.load(path, map_location=DEVICE))
-    return m.eval()
+# Define your model directly inside this file (to avoid import errors)
+import torch.nn as nn
 
-try:
-    tumor_model  = load_model(BrainTumorModel, BTD_FILENAME)
-    glioma_model = load_model(GliomaStageModel, GLIO_FILENAME)
-except Exception as e:
-    print("❌ Error loading model:", e)
+class BrainTumorModel(nn.Module):
+    def __init__(self):
+        super(BrainTumorModel, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, kernel_size=3),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Flatten(),
+            nn.Linear(32 * 54 * 54, 2),
+        )
 
-@app.get("/")
-async def health():
-    return {"status": "ok", "message": "API is up"}
+    def forward(self, x):
+        return self.model(x)
 
-@app.post("/predict-image/")
-async def predict_image(file: UploadFile = File(...)):
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(400, "Upload an image")
-    img = Image.open(file.file).convert("RGB")
-    t   = transform(img).unsqueeze(0).to(DEVICE)
-    with torch.no_grad():
-        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"}
-    return {
-        "tumor_type": tumor_type,
-        "precaution": get_precautions_from_gemini(tumor_type)
-    }
+# Load model
+model_path = "BTD_model.pth"
+if not os.path.exists(model_path):
+    from huggingface_hub import hf_hub_download
+    model_path = hf_hub_download(repo_id="Codewithsalty/brain-tumor-models", filename="BTD_model.pth", cache_dir="/tmp/.cache")
 
-class MutationInput(BaseModel):
-    gender: str
-    age:    float
-    idh1:   int
-    tp53:   int
-    atrx:   int
-    pten:   int
-    egfr:   int
-    cic:    int
-    pik3ca: int
+btd_model = BrainTumorModel()
+btd_model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+btd_model.eval()
 
-@app.post("/predict-glioma-stage/")
-async def predict_glioma_stage(data: MutationInput):
-    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(t)
-        idx = int(out.argmax(1))
-    stages = ["Stage 1","Stage 2","Stage 3","Stage 4"]
-    return {"glioma_stage": stages[idx]}
+# Define prediction endpoint
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        image_tensor = transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            output = btd_model(image_tensor)
+            prediction = torch.argmax(output, dim=1).item()
+
+        result = {0: "No tumor", 1: "Tumor detected"}[prediction]
+        return {"prediction": result}
+
+    except Exception as e:
+        return {"error": str(e)}
+
+# Health check endpoint
+@app.get("/")
+def root():
+    return {"message": "🧠 Brain Tumor Detection API is running!"}