Updating model

api.py

@@ -18,9 +18,9 @@ from utils import unzip_dataset, clean_dir
 
 app = FastAPI(title="RETFound MAE – Train & Inference API")
 
-# ---------- Config (env vars for Spaces) ----------
-DATA_ROOT = os.getenv("DATA_ROOT", "./workspace/data")
-CKPT_DIR  = os.getenv("CKPT_DIR",  "./workspace/checkpoints")
+# ---------- Config (safe paths for HF Spaces) ----------
+DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")
+CKPT_DIR  = os.getenv("CKPT_DIR",  "/tmp/checkpoints")
 BASE_REPO = os.getenv("HF_BASE_MODEL_REPO", "YukunZhou/RETFound_mae_meh")
 BASE_FILE = os.getenv("HF_BASE_MODEL_FILE", "RETFound_mae_meh.pth")
 MODEL_PUSH_REPO = os.getenv("HF_PUSH_REPO", "habeebCycle/RETFound_mae_meh_1")
@@ -44,7 +44,6 @@ _transform = transforms.Compose([
 
 def _load_model_for_inference():
     global _model
-    # If we already trained and have a best checkpoint, load it; else build unfitted head
     if _state["best_ckpt"] and os.path.exists(_state["best_ckpt"]):
         ckpt = torch.load(_state["best_ckpt"], map_location=DEVICE)
         classes = ckpt.get("classes", [])
@@ -52,7 +51,6 @@ def _load_model_for_inference():
         model = build_classifier(num_classes=len(classes), base_repo=BASE_REPO, base_filename=BASE_FILE, device=DEVICE)
         model.load_state_dict(ckpt["model"], strict=False)
     else:
-        # Cold start with 2 placeholder classes to allow /predict errors to be informative
         model = build_classifier(num_classes=2, base_repo=BASE_REPO, base_filename=BASE_FILE, device=DEVICE)
     model.eval()
     _model = model
@@ -71,17 +69,13 @@ def status():
 
 @app.post("/upload_dataset")
 async def upload_dataset(file: UploadFile = File(...)):
-    """Upload a ZIP that contains train/ and val/ folders.
-    Example structure inside zip:
-    train/class_a/*.jpg, val/class_a/*.jpg
-    """
+    """Upload a ZIP that contains train/ and val/ folders."""
     os.makedirs("/tmp/uploads", exist_ok=True)
     zip_path = f"/tmp/uploads/{file.filename}"
     with open(zip_path, "wb") as f:
         f.write(await file.read())
 
-    # Clean previous dataset and extract
-    clean_dir(DATA_ROOT)
+    clean_dir(DATA_ROOT)  # Now points to /tmp/data
     extracted = unzip_dataset(zip_path, DATA_ROOT)
     return {"ok": True, "dataset_dir": extracted}
 
@@ -141,9 +135,7 @@ async def predict(file: UploadFile = File(...)):
 
 @app.post("/push")
 async def push_to_hub(repo_id: Optional[str] = Form(None)):
-    """Push best checkpoint + metadata to Hugging Face Hub model repo.
-    Requires write permission via HF token (Spaces secret HF_TOKEN is used automatically).
-    """
+    """Push best checkpoint + metadata to Hugging Face Hub."""
     repo_id = repo_id or MODEL_PUSH_REPO
     if not repo_id:
         return JSONResponse({"error": "Set HF_PUSH_REPO env var or pass repo_id."}, status_code=400)
@@ -153,10 +145,8 @@ async def push_to_hub(repo_id: Optional[str] = Form(None)):
     api = HfApi()
     create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
 
-    # Upload files
     upload_file(path_or_fileobj=_state["best_ckpt"], path_in_repo="retfound_classifier_best.pth", repo_id=repo_id, repo_type="model")
 
-    # Write simple model card
     card = f"""# RETFound MAE – Retinal Classifier (Fine-tuned)
 
 - Base: `{BASE_REPO}/{BASE_FILE}`
@@ -171,4 +161,4 @@ This repo contains a PyTorch checkpoint `retfound_classifier_best.pth` compatibl
         f.write(card)
     upload_file(path_or_fileobj=card_path, path_in_repo="README.md", repo_id=repo_id, repo_type="model")
 
-    return {"ok": True, "pushed_to": repo_id}
+    return {"ok": True, "pushed_to": repo_id}