Update newapi.py

newapi.py

@@ -1,72 +1,65 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import torch
-from torchvision import transforms
-from PIL import Image
-import io
 import os
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from utils import BrainTumorModel, GliomaStageModel
 
-# ✅ 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
-
-# Define your app
 app = FastAPI()
 
-# ✅ Enable CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Load models (updated to local .pth files)
+btd_model_path = "brain_tumor_model.pth"
+glioma_model_path = "glioma_stage_model.pth"
 
-# ✅ 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")
+# Initialize and load Brain Tumor Detection Model
+btd_model = BrainTumorModel()
+btd_model.load_state_dict(torch.load(btd_model_path, map_location=torch.device('cpu')))
+btd_model.eval()
 
-btd_model = TumorClassification(model_path=btd_model_path)
-glioma_model = GliomaStageModel(model_path=glioma_model_path)
+# Initialize and load Glioma Stage Detection Model
+glioma_model = GliomaStageModel()
+glioma_model.load_state_dict(torch.load(glioma_model_path, map_location=torch.device('cpu')))
+glioma_model.eval()
 
-# ✅ Image preprocessing
+# Define preprocessing
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
 ])
 
-class DiagnosisResponse(BaseModel):
-    tumor: str
-    stage: str
-    precautions: list
-
-@app.post("/predict", response_model=DiagnosisResponse)
+@app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    image_tensor = transform(image).unsqueeze(0)
+    try:
+        image = Image.open(file.file).convert("RGB")
+        image = transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            output = btd_model(image)
+            predicted = torch.argmax(output, dim=1).item()
+
+        classes = ['No Tumor', 'Pituitary', 'Meningioma', 'Glioma']
+        result = classes[predicted]
+
+        return JSONResponse(content={"prediction": result})
+
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)})
+
+@app.post("/glioma-stage/")
+async def glioma_stage(file: UploadFile = File(...)):
+    try:
+        image = Image.open(file.file).convert("RGB")
+        image = 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=[]
-        )
+        with torch.no_grad():
+            output = glioma_model(image)
+            predicted = torch.argmax(output, dim=1).item()
 
-    stage_result = glioma_model.predict(image_tensor)
-    precautions = get_precautions_from_gemini(tumor_result, stage_result)
+        stages = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
+        result = stages[predicted]
 
-    return DiagnosisResponse(
-        tumor=tumor_result,
-        stage=stage_result,
-        precautions=precautions
-    )
+        return JSONResponse(content={"glioma_stage": result})
 
-@app.get("/")
-def root():
-    return {"message": "Brain Tumor API is running."}
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)})