Update newapi.py

newapi.py

@@ -6,17 +6,18 @@ from torchvision import transforms
 from PIL import Image
 import io
 import os
-from huggingface_hub import hf_hub_download
 
+# ✅ Set Hugging Face model cache directory to a writable path
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+
+from huggingface_hub import hf_hub_download
 from models.TumorModel import TumorClassification, GliomaStageModel
 from utils import get_precautions_from_gemini
 
-# ✅ Use /data as Hugging Face allows writing here only
-cache_dir = os.path.join("/data", "cache")
-os.makedirs(cache_dir, exist_ok=True)
-
-app = FastAPI(title="Brain Tumor Detection API")
+# Define your app
+app = FastAPI()
 
+# ✅ Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -25,83 +26,47 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ✅ Load Tumor Classification Model
-btd_model_path = hf_hub_download(
-    repo_id="Codewithsalty/brain-tumor-models",
-    filename="BTD_model.pth",
-    cache_dir=cache_dir
-)
-tumor_model = TumorClassification()
-tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
-tumor_model.eval()
+# ✅ Load your models from the Hugging Face Hub
+btd_model_path = hf_hub_download(repo_id="Codewithsalty/brain-tumor-detection", filename="brain_tumor_model.pt")
+glioma_model_path = hf_hub_download(repo_id="Codewithsalty/brain-tumor-detection", filename="glioma_stage_model.pt")
 
-# ✅ Load Glioma Stage Model
-glioma_model_path = hf_hub_download(
-    repo_id="Codewithsalty/brain-tumor-models",
-    filename="glioma_stages.pth",
-    cache_dir=cache_dir
-)
-glioma_model = GliomaStageModel()
-glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))
-glioma_model.eval()
+btd_model = TumorClassification(model_path=btd_model_path)
+glioma_model = GliomaStageModel(model_path=glioma_model_path)
 
+# ✅ Image preprocessing
 transform = transforms.Compose([
-    transforms.Grayscale(),
     transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5], std=[0.5]),
+    transforms.ToTensor()
 ])
 
-@app.get("/")
-async def root():
-    return {"message": "Brain Tumor Detection API is running."}
-
-labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
+class DiagnosisResponse(BaseModel):
+    tumor: str
+    stage: str
+    precautions: list
+
+@app.post("/predict", response_model=DiagnosisResponse)
+async def predict(file: UploadFile = File(...)):
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents)).convert("RGB")
+    image_tensor = transform(image).unsqueeze(0)
+
+    tumor_result = btd_model.predict(image_tensor)
+    if tumor_result == "No Tumor":
+        return DiagnosisResponse(
+            tumor="No Tumor Detected",
+            stage="N/A",
+            precautions=[]
+        )
+
+    stage_result = glioma_model.predict(image_tensor)
+    precautions = get_precautions_from_gemini(tumor_result, stage_result)
+
+    return DiagnosisResponse(
+        tumor=tumor_result,
+        stage=stage_result,
+        precautions=precautions
+    )
 
-@app.post("/predict-image")
-async def predict_image(file: UploadFile = File(...)):
-    img_bytes = await file.read()
-    img = Image.open(io.BytesIO(img_bytes)).convert("L")
-    x = transform(img).unsqueeze(0)
-
-    with torch.no_grad():
-        out = tumor_model(x)
-        idx = torch.argmax(out, dim=1).item()
-        tumor_type = labels[idx]
-
-    if tumor_type == "glioma":
-        return {"tumor_type": tumor_type, "next": "submit_mutation_data"}
-    else:
-        precautions = get_precautions_from_gemini(tumor_type)
-        return {"tumor_type": tumor_type, "precaution": precautions}
-
-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() == 'm' else 1
-    features = [
-        gender_val, data.age, data.idh1, data.tp53, data.atrx,
-        data.pten, data.egfr, data.cic, data.pik3ca
-    ]
-    x = torch.tensor(features).float().unsqueeze(0)
-
-    with torch.no_grad():
-        out = glioma_model(x)
-        idx = torch.argmax(out, dim=1).item()
-        stages = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
-        return {"glioma_stage": stages[idx]}
-
-# Only used when running locally
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("newapi:app", host="0.0.0.0", port=10000)
+@app.get("/")
+def root():
+    return {"message": "Brain Tumor API is running."}