Update newapi.py

newapi.py

@@ -5,24 +5,33 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 from utils import BrainTumorModel, GliomaStageModel
+from huggingface_hub import hf_hub_download
 
 app = FastAPI()
 
-# Load models (updated to local .pth files)
-btd_model_path = "brain_tumor_model.pth"
-glioma_model_path = "glioma_stage_model.pth"
+# Download models from the Space's repo
+btd_model_path = hf_hub_download(
+    repo_id="Codewithsalty/brain-tumor-api",
+    filename="brain_tumor_model.pth",
+    repo_type="space"
+)
 
-# Initialize and load Brain Tumor Detection Model
+glioma_model_path = hf_hub_download(
+    repo_id="Codewithsalty/brain-tumor-api",
+    filename="glioma_stage_model.pth",
+    repo_type="space"
+)
+
+# Load and prepare models
 btd_model = BrainTumorModel()
 btd_model.load_state_dict(torch.load(btd_model_path, map_location=torch.device('cpu')))
 btd_model.eval()
 
-# 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()
 
-# Define preprocessing
+# Image preprocessing
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -62,4 +71,4 @@ async def glioma_stage(file: UploadFile = File(...)):
         return JSONResponse(content={"glioma_stage": result})
 
     except Exception as e:
-        return JSONResponse(content={"error": str(e)})
+        return JSONResponse(content={"error": str(e)})