Initial Space

backend/app.py

@@ -1,3 +1,22 @@
+import os
+import shutil
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface"
+os.environ["TORCH_HOME"] = "/tmp/torch"
+os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
+os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
+
+for d in ["/tmp/huggingface", "/tmp/Ultralytics", "/tmp/matplotlib", "/tmp/torch", "/root/.cache"]:
+    shutil.rmtree(d, ignore_errors=True)
+
+from huggingface_hub import login
+
+
+
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
@@ -15,12 +34,6 @@ from model import MWT as create_model
 from augmentations import Augmentations
 from model_histo import BreastCancerClassifier  # TensorFlow model
 
-from huggingface_hub import login
-import os
-
-hf_token = os.getenv("HF_TOKEN")
-if hf_token:
-    login(token=hf_token)
 
 
 # =====================================================
@@ -36,7 +49,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-OUTPUT_DIR = "outputs"
+OUTPUT_DIR = "/tmp/outputs"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")
 
@@ -55,7 +68,7 @@ print("🔹 Loading MWT model...")
 mwt_model = create_model(num_classes=2).to(device)
 mwt_model.load_state_dict(torch.load("MWTclass2.pth", map_location=device))
 mwt_model.eval()
-mwt_class_names = ['neg', 'pos']
+mwt_class_names = ['Negative', 'Positive']
 
 # =====================================================
 # Model 3: CIN Classifier