diff --git a/CorpusCallosum/localization/inference.py b/CorpusCallosum/localization/inference.py index 595e0d8f0..062fd01b9 100644 --- a/CorpusCallosum/localization/inference.py +++ b/CorpusCallosum/localization/inference.py @@ -22,10 +22,12 @@ from CorpusCallosum.transforms.localization import CropAroundACPCFixedSize from CorpusCallosum.utils.types import Points2dType -from FastSurferCNN.download_checkpoints import load_checkpoint_config_defaults -from FastSurferCNN.download_checkpoints import main as download_checkpoints from FastSurferCNN.utils import Image3d, Vector2d, Vector3d -from FastSurferCNN.utils.checkpoint import get_config_file +from FastSurferCNN.utils.checkpoint import ( + get_checkpoints, + get_config_file, + load_checkpoint_config_defaults, +) from FastSurferCNN.utils.parser_defaults import FASTSURFER_ROOT PATCH_SIZE = (64, 64) @@ -59,10 +61,11 @@ def load_model(device: torch.device) -> DenseNet: dropout_prob=0.2 ) - download_checkpoints(cc=True) config_file = get_config_file("CorpusCallosum") cc_config = load_checkpoint_config_defaults("checkpoint", filename=config_file) + urls = load_checkpoint_config_defaults("url", filename=config_file) checkpoint_path = FASTSURFER_ROOT / cc_config['localization'] + get_checkpoints(checkpoint_path, urls=urls) # Load state dict if isinstance(checkpoint_path, str) or isinstance(checkpoint_path, Path): diff --git a/CorpusCallosum/segmentation/inference.py b/CorpusCallosum/segmentation/inference.py index 2e4a384e4..640972d50 100644 --- a/CorpusCallosum/segmentation/inference.py +++ b/CorpusCallosum/segmentation/inference.py @@ -23,11 +23,13 @@ from CorpusCallosum.data import constants from CorpusCallosum.transforms.segmentation import CropAroundACPC -from FastSurferCNN.download_checkpoints import load_checkpoint_config_defaults -from FastSurferCNN.download_checkpoints import main as download_checkpoints from FastSurferCNN.models.networks import FastSurferVINN from FastSurferCNN.utils import Image3d, Image4d, Shape2d, Shape3d, Shape4d, Vector2d, nibabelImage -from FastSurferCNN.utils.checkpoint import get_config_file +from FastSurferCNN.utils.checkpoint import ( + get_checkpoints, + get_config_file, + load_checkpoint_config_defaults, +) from FastSurferCNN.utils.parallel import thread_executor @@ -69,10 +71,11 @@ def load_model(device: torch.device | None = None) -> FastSurferVINN: } model = FastSurferVINN(params) - download_checkpoints(cc=True) config_file = get_config_file("CorpusCallosum") cc_config: dict[str, Path] = load_checkpoint_config_defaults("checkpoint", filename=config_file) + urls = load_checkpoint_config_defaults("url", filename=config_file) checkpoint_path = constants.FASTSURFER_ROOT / cc_config['segmentation'] + get_checkpoints(checkpoint_path, urls=urls) weights = torch.load(checkpoint_path, weights_only=True, map_location=device) model.load_state_dict(weights) diff --git a/FastSurferCNN/utils/checkpoint.py b/FastSurferCNN/utils/checkpoint.py index 2c4e47fe8..701aaeaaa 100644 --- a/FastSurferCNN/utils/checkpoint.py +++ b/FastSurferCNN/utils/checkpoint.py @@ -19,6 +19,7 @@ from functools import lru_cache from pathlib import Path from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload +from uuid import uuid4 import requests import torch @@ -400,8 +401,20 @@ def download_checkpoint( raise RuntimeError(message, responses) else: response = next(r for r in responses if r.ok) - with open(checkpoint_path, "wb") as f: - f.write(response.content) + checkpoint_path = Path(checkpoint_path) + temporary_path = checkpoint_path.with_name( + f".{checkpoint_path.name}.{uuid4().hex}.tmp" + ) + try: + # Opening a unique file with ``xb`` preserves the permissions dictated by + # the process umask and prevents concurrent downloads from sharing a + # temporary file. The same-directory replace atomically publishes the + # checkpoint only after it has been written and closed completely. + with open(temporary_path, "xb") as f: + f.write(response.content) + os.replace(temporary_path, checkpoint_path) + finally: + temporary_path.unlink(missing_ok=True) def check_and_download_ckpts(checkpoint_path: Path | str, urls: list[str]) -> None: