| """ |
| Hugging Face dataset sync. |
| |
| Source dataset layout (CTSpinoPelvic1K, flat): |
| ct/{case_id}_ct.nii.gz |
| labels/{case_id}_label.nii.gz |
| |
| where case_id follows export_hf.py's filename convention: |
| {token:04d}_{position} # fused |
| {token:04d}_{position}_spine # separate, spine view |
| {token:04d}_{position}_pelvic # separate, pelvic view |
| e.g. "0189_unknown", "0401_unknown_pelvic" |
| |
| Target dataset layout (CTSpinoPelvic1K-annotations): |
| labels/{case_id}/{username}_{ISO8601}.nii.gz # versioned refinements |
| audit.sqlite # provenance DB |
| |
| The source label is offered to OHIF as a starting segmentation so clinicians |
| refine rather than redraw. This matches the paper's thesis: TotalSegmentator's |
| automated labels have a 19.6 junction-DSC gap on LSTV cases, so the value of |
| the annotation pass is corrective. |
| """ |
| from __future__ import annotations |
|
|
| import asyncio |
| import logging |
| import re |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
| from huggingface_hub import HfApi, hf_hub_download |
| from huggingface_hub.utils import EntryNotFoundError |
|
|
| log = logging.getLogger("sync_manager") |
|
|
| _SAFE = re.compile(r"[^A-Za-z0-9_\-]") |
|
|
|
|
| def _sanitize(s: str) -> str: |
| return _SAFE.sub("_", s) |
|
|
|
|
| class SyncManager: |
| def __init__( |
| self, |
| source_dataset: str, |
| target_dataset: str, |
| workspace: str, |
| hf_token: str, |
| audit_logger: Any, |
| ): |
| self.source = source_dataset |
| self.target = target_dataset |
| self.workspace = Path(workspace) |
| self.raw = self.workspace / "raw_data" |
| self.labels = self.workspace / "labels" |
| self.audit_db = self.workspace / "audit.sqlite" |
|
|
| self.raw.mkdir(parents=True, exist_ok=True) |
| self.labels.mkdir(parents=True, exist_ok=True) |
|
|
| self.token = hf_token |
| self.api = HfApi(token=hf_token) |
| self.audit = audit_logger |
|
|
| |
| self._upload_lock = asyncio.Lock() |
|
|
| |
| |
| |
|
|
| def ensure_case(self, case_id: str) -> dict: |
| """ |
| Ensure CT + source label for a case are present locally. Idempotent. |
| |
| Returns paths so the caller (FastAPI) can pass them to MONAI Label's |
| datastore for registration. The source label is the *starting* mask — |
| OHIF loads it and the annotator refines on top. |
| """ |
| safe = _sanitize(case_id) |
| case_dir = self.raw / safe |
| case_dir.mkdir(parents=True, exist_ok=True) |
|
|
| local_ct = case_dir / "ct.nii.gz" |
| local_seed_label = case_dir / "seed_label.nii.gz" |
|
|
| ct_remote = f"ct/{case_id}_ct.nii.gz" |
| label_remote = f"labels/{case_id}_label.nii.gz" |
|
|
| if not local_ct.exists(): |
| try: |
| downloaded = hf_hub_download( |
| repo_id=self.source, |
| filename=ct_remote, |
| repo_type="dataset", |
| token=self.token, |
| ) |
| self._link_or_copy(Path(downloaded), local_ct) |
| except EntryNotFoundError as e: |
| raise FileNotFoundError( |
| f"CT not found for case {case_id!r} at {ct_remote} " |
| f"in {self.source}: {e}" |
| ) from e |
|
|
| |
| if not local_seed_label.exists(): |
| try: |
| downloaded = hf_hub_download( |
| repo_id=self.source, |
| filename=label_remote, |
| repo_type="dataset", |
| token=self.token, |
| ) |
| self._link_or_copy(Path(downloaded), local_seed_label) |
| except EntryNotFoundError: |
| log.info("No seed label for %s; annotator starts from blank", case_id) |
|
|
| return { |
| "ct": str(local_ct), |
| "seed_label": str(local_seed_label) if local_seed_label.exists() else None, |
| "case_dir": str(case_dir), |
| } |
|
|
| @staticmethod |
| def _link_or_copy(src: Path, dst: Path) -> None: |
| try: |
| dst.symlink_to(src) |
| except OSError: |
| dst.write_bytes(src.read_bytes()) |
|
|
| |
| |
| |
|
|
| def _versioned_mask_path(self, case_id: str, username: str) -> Path: |
| """labels/<case>/<user>_<ISO8601>.nii.gz — concurrent-safe.""" |
| ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") |
| case_dir = self.labels / _sanitize(case_id) |
| case_dir.mkdir(parents=True, exist_ok=True) |
| return case_dir / f"{_sanitize(username)}_{ts}.nii.gz" |
|
|
| async def save_mask( |
| self, |
| case_id: str, |
| username: str, |
| mask_bytes: bytes, |
| session_id: str, |
| ) -> dict: |
| """Persist refined mask locally, log the event, push to target dataset.""" |
| out = self._versioned_mask_path(case_id, username) |
| await asyncio.to_thread(out.write_bytes, mask_bytes) |
|
|
| sha = self.audit.log_annotation( |
| username=username, |
| case_id=case_id, |
| source_filename=f"ct/{case_id}_ct.nii.gz", |
| mask_filename=str(out.relative_to(self.workspace)), |
| mask_bytes=mask_bytes, |
| session_id=session_id, |
| ) |
|
|
| |
| asyncio.create_task(self._push_artifacts(out, sha, case_id, username)) |
|
|
| return { |
| "path": str(out.relative_to(self.workspace)), |
| "sha256": sha, |
| "bytes": len(mask_bytes), |
| } |
|
|
| async def _push_artifacts( |
| self, mask_path: Path, sha: str, case_id: str, username: str, |
| ) -> None: |
| """Upload the refined mask and the updated audit DB to target.""" |
| async with self._upload_lock: |
| try: |
| await asyncio.to_thread( |
| self.api.upload_file, |
| path_or_fileobj=str(mask_path), |
| path_in_repo=str(mask_path.relative_to(self.workspace)), |
| repo_id=self.target, |
| repo_type="dataset", |
| commit_message=( |
| f"annotation: {case_id} by {username} " |
| f"(sha256={sha[:12]})" |
| ), |
| ) |
| if self.audit_db.exists(): |
| await asyncio.to_thread( |
| self.api.upload_file, |
| path_or_fileobj=str(self.audit_db), |
| path_in_repo="audit.sqlite", |
| repo_id=self.target, |
| repo_type="dataset", |
| commit_message=f"audit: after {mask_path.name}", |
| ) |
| log.info("Pushed %s to %s", mask_path.name, self.target) |
| except Exception: |
| log.exception("Failed to push %s to %s", mask_path, self.target) |
|
|
