Update newapi.py

newapi.py

@@ -2,21 +2,20 @@
 
 from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
 import io
 import os
 
-# Use a writable directory in Hugging Face Spaces
+# Set writable cache directories
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache"
 os.environ["HF_HOME"] = "/tmp/.cache"
 
-# Define FastAPI app
+# FastAPI setup
 app = FastAPI(title="🧠 Brain Tumor Detection API")
 
-# Enable CORS
+# Allow CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -24,26 +23,27 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Image transform
+# Define image transform (grayscale)
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
+    transforms.Grayscale(num_output_channels=1),  # Ensure grayscale
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5], std=[0.5]),
 ])
 
-# Define your model directly inside this file (to avoid import errors)
+# Define the exact same model used during training
 import torch.nn as nn
 
 class BrainTumorModel(nn.Module):
     def __init__(self):
         super(BrainTumorModel, self).__init__()
-        self.con1d = nn.Conv2d(3, 32, kernel_size=3)
+        self.con1d = nn.Conv2d(1, 32, kernel_size=3)   # Input is grayscale (1 channel)
         self.con2d = nn.Conv2d(32, 64, kernel_size=3)
         self.con3d = nn.Conv2d(64, 128, kernel_size=3)
         self.pool = nn.MaxPool2d(2)
-        self.fc1 = nn.Linear(128 * 25 * 25, 256)
-        self.fc2 = nn.Linear(256, 128)
-        self.output = nn.Linear(128, 2)
+        self.fc1 = nn.Linear(128 * 25 * 25, 512)      # Matches your saved model
+        self.fc2 = nn.Linear(512, 256)
+        self.output = nn.Linear(256, 4)               # 4 classes expected
 
     def forward(self, x):
         x = self.pool(torch.relu(self.con1d(x)))
@@ -65,25 +65,31 @@ btd_model = BrainTumorModel()
 btd_model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 btd_model.eval()
 
-# Define prediction endpoint
+# 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 = Image.open(io.BytesIO(contents)).convert("L")  # Grayscale
         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]
+        result = {
+            0: "No tumor",
+            1: "Glioma",
+            2: "Meningioma",
+            3: "Pituitary tumor"
+        }[prediction]
+
         return {"prediction": result}
 
     except Exception as e:
         return {"error": str(e)}
 
-# Health check endpoint
+# Health check
 @app.get("/")
 def root():
     return {"message": "🧠 Brain Tumor Detection API is running!"}