Codewithsalty commited on
Commit
daee81e
·
verified ·
1 Parent(s): e4acaca

Update newapi.py

Browse files
Files changed (1) hide show
  1. newapi.py +9 -3
newapi.py CHANGED
@@ -5,11 +5,15 @@ import torch
5
  from torchvision import transforms
6
  from PIL import Image
7
  import io
 
8
  from huggingface_hub import hf_hub_download
9
 
10
  from models.TumorModel import TumorClassification, GliomaStageModel
11
  from utils import get_precautions_from_gemini
12
 
 
 
 
13
  # Initialize FastAPI app
14
  app = FastAPI(title="Brain Tumor Detection API")
15
 
@@ -25,7 +29,8 @@ app.add_middleware(
25
  # Load Tumor Classification Model from Hugging Face
26
  btd_model_path = hf_hub_download(
27
  repo_id="Codewithsalty/brain-tumor-models",
28
- filename="BTD_model.pth"
 
29
  )
30
  tumor_model = TumorClassification()
31
  tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
@@ -34,7 +39,8 @@ tumor_model.eval()
34
  # Load Glioma Stage Prediction Model from Hugging Face
35
  glioma_model_path = hf_hub_download(
36
  repo_id="Codewithsalty/brain-tumor-models",
37
- filename="glioma_stages.pth"
 
38
  )
39
  glioma_model = GliomaStageModel()
40
  glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))
@@ -104,4 +110,4 @@ async def predict_glioma_stage(data: MutationInput):
104
 
105
  if __name__ == "__main__":
106
  import uvicorn
107
- uvicorn.run("newapi:app", host="0.0.0.0", port=10000)
 
5
  from torchvision import transforms
6
  from PIL import Image
7
  import io
8
+ import os
9
  from huggingface_hub import hf_hub_download
10
 
11
  from models.TumorModel import TumorClassification, GliomaStageModel
12
  from utils import get_precautions_from_gemini
13
 
14
+ # Create writable cache directory for Hugging Face
15
+ os.makedirs("./cache", exist_ok=True)
16
+
17
  # Initialize FastAPI app
18
  app = FastAPI(title="Brain Tumor Detection API")
19
 
 
29
  # Load Tumor Classification Model from Hugging Face
30
  btd_model_path = hf_hub_download(
31
  repo_id="Codewithsalty/brain-tumor-models",
32
+ filename="BTD_model.pth",
33
+ cache_dir="./cache" # ✅ fix
34
  )
35
  tumor_model = TumorClassification()
36
  tumor_model.load_state_dict(torch.load(btd_model_path, map_location="cpu"))
 
39
  # Load Glioma Stage Prediction Model from Hugging Face
40
  glioma_model_path = hf_hub_download(
41
  repo_id="Codewithsalty/brain-tumor-models",
42
+ filename="glioma_stages.pth",
43
+ cache_dir="./cache" # ✅ fix
44
  )
45
  glioma_model = GliomaStageModel()
46
  glioma_model.load_state_dict(torch.load(glioma_model_path, map_location="cpu"))
 
110
 
111
  if __name__ == "__main__":
112
  import uvicorn
113
+ uvicorn.run("newapi:app", host="0.0.0.0", port=10000)