Update newapi.py

newapi.py

@@ -1,59 +1,51 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
 import torch
 import torchvision.transforms as transforms
-from PIL import Image
-import io
-
-from utils import YourModelClass  # Make sure this matches your actual model class
+from utils import BrainTumorModel, get_precautions_from_gemini
 
 app = FastAPI()
 
-# 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()
-
-# Image transformation (adjust according to how your model was trained)
+# Load the model
+btd_model = BrainTumorModel()
+btd_model_path = "brain_tumor_model.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}")
+
+# Define image transform
 transform = transforms.Compose([
-    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
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
 ])
 
+# Class labels (adjust if your model uses different labels)
+classes = ['glioma', 'meningioma', 'notumor', 'pituitary']
+
 @app.get("/")
-def root():
-    return {"message": "Brain Tumor Detection API is up and running!"}
+def read_root():
+    return {"message": "Brain Tumor Detection API is running 🚀"}
 
-@app.post("/predict/")
+@app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     try:
-        # Read image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert('RGB')
-        image = transform(image).unsqueeze(0)
+        image = Image.open(file.file).convert("RGB")
+        image = transform(image).unsqueeze(0)  # Shape: [1, 3, 224, 224]
 
-        # Run model
         with torch.no_grad():
             outputs = btd_model(image)
-            _, predicted = torch.max(outputs, 1)
+            _, predicted = torch.max(outputs.data, 1)
+            predicted_class = classes[predicted.item()]
+            precautions = get_precautions_from_gemini(predicted_class)
 
-        # Class mapping (adjust according to your model's labels)
-        classes = ['No Tumor', 'Glioma', 'Meningioma', 'Pituitary']
-        prediction = classes[predicted.item()]
+        return JSONResponse(content={
+            "prediction": predicted_class,
+            "precautions": precautions
+        })
 
-        return JSONResponse({"prediction": prediction})
-    
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return JSONResponse(content={"error": str(e)}, status_code=500)