Update newapi.py

newapi.py

@@ -5,20 +5,17 @@ import torch
 from torchvision import transforms
 from PIL import Image
 import io
-import os
 from huggingface_hub import hf_hub_download
 
 from models.TumorModel import TumorClassification, GliomaStageModel
 from utils import get_precautions_from_gemini
 
-# ✅ Use a safe local cache dir
-cache_dir = "./hf_cache"
-os.makedirs(cache_dir, exist_ok=True)  # create if it doesn't exist
+# ✅ Let Hugging Face handle cache automatically — DO NOT manually create any folders
 
 # Initialize FastAPI app
 app = FastAPI(title="Brain Tumor Detection API")
 
-# Enable CORS
+# Enable CORS for all origins
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -27,27 +24,25 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load tumor classification model
+# ✅ Load Tumor Classification Model from Hugging Face
 btd_model_path = hf_hub_download(
     repo_id="Codewithsalty/brain-tumor-models",
-    filename="BTD_model.pth",
-    cache_dir=cache_dir
+    filename="BTD_model.pth"
 )
 tumor_model = TumorClassification()
 tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
 tumor_model.eval()
 
-# Load glioma stage model
+# ✅ Load Glioma Stage Prediction Model from Hugging Face
 glioma_model_path = hf_hub_download(
     repo_id="Codewithsalty/brain-tumor-models",
-    filename="glioma_stages.pth",
-    cache_dir=cache_dir
+    filename="glioma_stages.pth"
 )
 glioma_model = GliomaStageModel()
 glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))
 glioma_model.eval()
 
-# Image preprocessing
+# Image preprocessing pipeline
 transform = transforms.Compose([
     transforms.Grayscale(),
     transforms.Resize((224, 224)),
@@ -55,17 +50,19 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.5], std=[0.5]),
 ])
 
+# Health check endpoint
 @app.get("/")
 async def root():
     return {"message": "Brain Tumor Detection API is running."}
 
-# Labels
+# Tumor type labels
 labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
 
+# Predict tumor type from uploaded image
 @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")
+    img = Image.open(io.BytesIO(img_bytes)).convert("L")  # Ensure grayscale
     x = transform(img).unsqueeze(0)
 
     with torch.no_grad():
@@ -79,7 +76,7 @@ async def predict_image(file: UploadFile = File(...)):
         precautions = get_precautions_from_gemini(tumor_type)
         return {"tumor_type": tumor_type, "precaution": precautions}
 
-# Mutation input
+# Input model for glioma mutation data
 class MutationInput(BaseModel):
     gender: str
     age: float
@@ -91,6 +88,7 @@ class MutationInput(BaseModel):
     cic: int
     pik3ca: int
 
+# Predict glioma stage based on mutations
 @app.post("/predict-glioma-stage")
 async def predict_glioma_stage(data: MutationInput):
     gender_val = 0 if data.gender.lower() == 'm' else 1
@@ -106,7 +104,4 @@ async def predict_glioma_stage(data: MutationInput):
         stages = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
         return {"glioma_stage": stages[idx]}
 
-# Only needed for local development, not in Hugging Face Spaces
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("newapi:app", host="0.0.0.0", port=10000)
+# ✅ No need to run uvicorn manually in Hugging Face Spaces