Update newapi.py

newapi.py

@@ -1,68 +1,59 @@
-import os
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
-from PIL import Image
+from fastapi.middleware.cors import CORSMiddleware
 import torch
 import torchvision.transforms as transforms
-from utils import BrainTumorModel, GliomaStageModel
+from PIL import Image
+import io
 
-app = FastAPI()
+from utils import YourModelClass  # Make sure this matches your actual model class
 
-# === Use exact filenames from the Space directory ===
-btd_model_path = "brain_tumor_model.pth"
-glioma_model_path = "glioma_stage_model.pth"
+app = FastAPI()
 
-# === Load Brain Tumor Model ===
-btd_model = BrainTumorModel()
+# CORS Middleware (optional but good for frontend API usage)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load model
+btd_model_path = "models/BTD_model.pth"  # ✅ Correct filename and folder
+btd_model = YourModelClass()
 btd_model.load_state_dict(torch.load(btd_model_path, map_location=torch.device('cpu')))
 btd_model.eval()
 
-# === Load Glioma Stage Model ===
-glioma_model = GliomaStageModel()
-glioma_model.load_state_dict(torch.load(glioma_model_path, map_location=torch.device('cpu')))
-glioma_model.eval()
-
-# === Image Transform ===
+# Image transformation (adjust according to how your model was trained)
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
+    transforms.Resize((224, 224)),  # Adjust to your model's expected input size
     transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5], std=[0.5])  # Adjust for grayscale or RGB
 ])
 
-# === Routes ===
+@app.get("/")
+def root():
+    return {"message": "Brain Tumor Detection API is up and running!"}
 
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     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")
+        # Read image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert('RGB')
         image = transform(image).unsqueeze(0)
 
+        # Run model
         with torch.no_grad():
-            output = glioma_model(image)
-            predicted = torch.argmax(output, dim=1).item()
-
-        stages = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
-        result = stages[predicted]
+            outputs = btd_model(image)
+            _, predicted = torch.max(outputs, 1)
 
-        return JSONResponse(content={"glioma_stage": result})
+        # Class mapping (adjust according to your model's labels)
+        classes = ['No Tumor', 'Glioma', 'Meningioma', 'Pituitary']
+        prediction = classes[predicted.item()]
 
+        return JSONResponse({"prediction": prediction})
+    
     except Exception as e:
-        return JSONResponse(content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": str(e)})