gregoryschwingmdphd/spinesurg-ct-annotator

1"""2Hugging Face dataset sync.3 4Source dataset layout (CTSpinoPelvic1K, flat):5    ct/{case_id}_ct.nii.gz6    labels/{case_id}_label.nii.gz7 8where case_id follows export_hf.py's filename convention:9    {token:04d}_{position}                       # fused10    {token:04d}_{position}_spine                  # separate, spine view11    {token:04d}_{position}_pelvic                 # separate, pelvic view12    e.g. "0189_unknown", "0401_unknown_pelvic"13 14Target dataset layout (CTSpinoPelvic1K-annotations):15    labels/{case_id}/{username}_{ISO8601}.nii.gz  # versioned refinements16    audit.sqlite                                   # provenance DB17 18The source label is offered to OHIF as a starting segmentation so clinicians19refine rather than redraw. This matches the paper's thesis: TotalSegmentator's20automated labels have a 19.6 junction-DSC gap on LSTV cases, so the value of21the annotation pass is corrective.22"""23from __future__ import annotations24 25import asyncio26import logging27import re28from datetime import datetime, timezone29from pathlib import Path30from typing import Any31 32from huggingface_hub import HfApi, hf_hub_download33from huggingface_hub.utils import EntryNotFoundError34 35log = logging.getLogger("sync_manager")36 37_SAFE = re.compile(r"[^A-Za-z0-9_\-]")38 39 40def _sanitize(s: str) -> str:41    return _SAFE.sub("_", s)42 43 44class SyncManager:45    def __init__(46        self,47        source_dataset: str,48        target_dataset: str,49        workspace: str,50        hf_token: str,51        audit_logger: Any,52    ):53        self.source = source_dataset54        self.target = target_dataset55        self.workspace = Path(workspace)56        self.raw = self.workspace / "raw_data"57        self.labels = self.workspace / "labels"58        self.audit_db = self.workspace / "audit.sqlite"59 60        self.raw.mkdir(parents=True, exist_ok=True)61        self.labels.mkdir(parents=True, exist_ok=True)62 63        self.token = hf_token64        self.api = HfApi(token=hf_token)65        self.audit = audit_logger66 67        # Serialize uploads to avoid rate-limit thrashing from concurrent saves.68        self._upload_lock = asyncio.Lock()69 70    # --------------------------------------------------------------------- #71    # Source dataset — lazy per-case download                               #72    # --------------------------------------------------------------------- #73 74    def ensure_case(self, case_id: str) -> dict:75        """76        Ensure CT + source label for a case are present locally. Idempotent.77 78        Returns paths so the caller (FastAPI) can pass them to MONAI Label's79        datastore for registration. The source label is the *starting* mask —80        OHIF loads it and the annotator refines on top.81        """82        safe = _sanitize(case_id)83        case_dir = self.raw / safe84        case_dir.mkdir(parents=True, exist_ok=True)85 86        local_ct = case_dir / "ct.nii.gz"87        local_seed_label = case_dir / "seed_label.nii.gz"88 89        ct_remote = f"ct/{case_id}_ct.nii.gz"90        label_remote = f"labels/{case_id}_label.nii.gz"91 92        if not local_ct.exists():93            try:94                downloaded = hf_hub_download(95                    repo_id=self.source,96                    filename=ct_remote,97                    repo_type="dataset",98                    token=self.token,99                )100                self._link_or_copy(Path(downloaded), local_ct)101            except EntryNotFoundError as e:102                raise FileNotFoundError(103                    f"CT not found for case {case_id!r} at {ct_remote} "104                    f"in {self.source}: {e}"105                ) from e106 107        # Seed label is optional — future datasets may not ship one.108        if not local_seed_label.exists():109            try:110                downloaded = hf_hub_download(111                    repo_id=self.source,112                    filename=label_remote,113                    repo_type="dataset",114                    token=self.token,115                )116                self._link_or_copy(Path(downloaded), local_seed_label)117            except EntryNotFoundError:118                log.info("No seed label for %s; annotator starts from blank", case_id)119 120        return {121            "ct": str(local_ct),122            "seed_label": str(local_seed_label) if local_seed_label.exists() else None,123            "case_dir": str(case_dir),124        }125 126    @staticmethod127    def _link_or_copy(src: Path, dst: Path) -> None:128        try:129            dst.symlink_to(src)130        except OSError:131            dst.write_bytes(src.read_bytes())132 133    # --------------------------------------------------------------------- #134    # Target dataset — versioned saves (option 3)                           #135    # --------------------------------------------------------------------- #136 137    def _versioned_mask_path(self, case_id: str, username: str) -> Path:138        """labels/<case>/<user>_<ISO8601>.nii.gz — concurrent-safe."""139        ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")140        case_dir = self.labels / _sanitize(case_id)141        case_dir.mkdir(parents=True, exist_ok=True)142        return case_dir / f"{_sanitize(username)}_{ts}.nii.gz"143 144    async def save_mask(145        self,146        case_id: str,147        username: str,148        mask_bytes: bytes,149        session_id: str,150    ) -> dict:151        """Persist refined mask locally, log the event, push to target dataset."""152        out = self._versioned_mask_path(case_id, username)153        await asyncio.to_thread(out.write_bytes, mask_bytes)154 155        sha = self.audit.log_annotation(156            username=username,157            case_id=case_id,158            source_filename=f"ct/{case_id}_ct.nii.gz",159            mask_filename=str(out.relative_to(self.workspace)),160            mask_bytes=mask_bytes,161            session_id=session_id,162        )163 164        # Fire-and-forget — OHIF shouldn't wait on HF's response time.165        asyncio.create_task(self._push_artifacts(out, sha, case_id, username))166 167        return {168            "path": str(out.relative_to(self.workspace)),169            "sha256": sha,170            "bytes": len(mask_bytes),171        }172 173    async def _push_artifacts(174        self, mask_path: Path, sha: str, case_id: str, username: str,175    ) -> None:176        """Upload the refined mask and the updated audit DB to target."""177        async with self._upload_lock:178            try:179                await asyncio.to_thread(180                    self.api.upload_file,181                    path_or_fileobj=str(mask_path),182                    path_in_repo=str(mask_path.relative_to(self.workspace)),183                    repo_id=self.target,184                    repo_type="dataset",185                    commit_message=(186                        f"annotation: {case_id} by {username} "187                        f"(sha256={sha[:12]})"188                    ),189                )190                if self.audit_db.exists():191                    await asyncio.to_thread(192                        self.api.upload_file,193                        path_or_fileobj=str(self.audit_db),194                        path_in_repo="audit.sqlite",195                        repo_id=self.target,196                        repo_type="dataset",197                        commit_message=f"audit: after {mask_path.name}",198                    )199                log.info("Pushed %s to %s", mask_path.name, self.target)200            except Exception:201                log.exception("Failed to push %s to %s", mask_path, self.target)202