From f65ff0fa594c3d3da77f24aa0b42e18c82cd286f Mon Sep 17 00:00:00 2001 From: Arvin Razavi Date: Tue, 14 Jul 2026 17:10:02 +0200 Subject: [PATCH 1/3] Write downloaded checkpoints atomically to avoid a load race download_checkpoint() wrote the HTTP response straight to the final path, so a concurrent caller could observe the file through the exists() guard in check_and_download_ckpts() while it was still being written and torch.load a truncated checkpoint. This reliably crashes the CorpusCallosum module on a fresh run, where the localization and segmentation models are loaded in parallel threads that both call download_checkpoints(): RuntimeError: PytorchStreamReader failed reading file data/1: file read failed Download to a temporary file in the same directory and os.replace() it into place, so the checkpoint only ever becomes visible at its final path once it is complete. --- FastSurferCNN/utils/checkpoint.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/FastSurferCNN/utils/checkpoint.py b/FastSurferCNN/utils/checkpoint.py index 2c4e47fe8..39176eb74 100644 --- a/FastSurferCNN/utils/checkpoint.py +++ b/FastSurferCNN/utils/checkpoint.py @@ -15,6 +15,7 @@ # IMPORTS import os import sys +import tempfile from collections.abc import MutableSequence from functools import lru_cache from pathlib import Path @@ -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) + # Write atomically: download to a temporary file in the same directory, then + # os.replace() it into place. This prevents a concurrent caller (e.g. the CC + # module loads its localization and segmentation models in parallel threads, + # both triggering download_checkpoints) from seeing a half-written file via the + # `checkpoint_path.exists()` guard and torch.load-ing a truncated checkpoint. + checkpoint_path = Path(checkpoint_path) + fd, tmp_path = tempfile.mkstemp(dir=checkpoint_path.parent, suffix=".part") + try: + with os.fdopen(fd, "wb") as f: + f.write(response.content) + os.replace(tmp_path, checkpoint_path) + except BaseException: + Path(tmp_path).unlink(missing_ok=True) + raise def check_and_download_ckpts(checkpoint_path: Path | str, urls: list[str]) -> None: From 5c0a22231f14e500a08b68a664c95baf9ab78f7b Mon Sep 17 00:00:00 2001 From: Arvin Razavi Date: Tue, 14 Jul 2026 17:10:02 +0200 Subject: [PATCH 2/3] Fix corpus callosum measurements under NumPy >= 2 NumPy 2.0 removed the 2-element form of np.cross, which now raises "Both input arrays must be (arrays of) 3-dimensional vectors". The corpus callosum shape code used np.cross only to obtain the sign (z-component) of the angle between two in-slice 2D vectors, so every slice's measurements failed: ValueError: Both input arrays must be (arrays of) 3-dimensional vectors, but they are 2 and 2 dimensional instead. Compute the 2D cross-product scalar directly (a[0]*b[1] - a[1]*b[0]). This is equivalent and works on the pinned numpy (2.4.4) as well as newer releases. --- CorpusCallosum/shape/subsegment_contour.py | 5 +++-- CorpusCallosum/utils/mapping_helpers.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/CorpusCallosum/shape/subsegment_contour.py b/CorpusCallosum/shape/subsegment_contour.py index 2434081df..801489c00 100644 --- a/CorpusCallosum/shape/subsegment_contour.py +++ b/CorpusCallosum/shape/subsegment_contour.py @@ -880,8 +880,9 @@ def transform_to_acpc_standard( norms_product = np.linalg.norm(ac_pc_vec) * np.linalg.norm(posterior_vector) theta = np.arccos(dot_product / norms_product) - # Determine the sign of the angle using cross product - cross_product = np.cross(ac_pc_vec, posterior_vector) + # Determine the sign of the angle using the 2D cross product (z-component). + # NumPy 2.0 removed np.cross for 2-element vectors, so compute the scalar directly. + cross_product = ac_pc_vec[0] * posterior_vector[1] - ac_pc_vec[1] * posterior_vector[0] if cross_product < 0: theta = -theta diff --git a/CorpusCallosum/utils/mapping_helpers.py b/CorpusCallosum/utils/mapping_helpers.py index 666e5e08b..b02dbee05 100644 --- a/CorpusCallosum/utils/mapping_helpers.py +++ b/CorpusCallosum/utils/mapping_helpers.py @@ -84,8 +84,9 @@ def correct_nodding(ac_pt: Vector2d, pc_pt: Vector2d) -> AffineMatrix3x3: norms_product = np.linalg.norm(ac_pc_vec) * np.linalg.norm(posterior_vector) theta = np.arccos(dot_product / norms_product) - # Determine the sign of the angle using cross product - cross_product = np.cross(ac_pc_vec, posterior_vector) + # Determine the sign of the angle using the 2D cross product (z-component). + # NumPy 2.0 removed np.cross for 2-element vectors, so compute the scalar directly. + cross_product = ac_pc_vec[0] * posterior_vector[1] - ac_pc_vec[1] * posterior_vector[0] if cross_product < 0: theta = -theta From 3ad1adf4d2715c3c96e60645de07da98f8778c9f Mon Sep 17 00:00:00 2001 From: ClePol Date: Tue, 21 Jul 2026 18:24:56 +0200 Subject: [PATCH 3/3] Download checkpoints atomically per model --- CorpusCallosum/localization/inference.py | 11 ++++++---- CorpusCallosum/segmentation/inference.py | 11 ++++++---- CorpusCallosum/shape/subsegment_contour.py | 5 ++--- CorpusCallosum/utils/mapping_helpers.py | 5 ++--- FastSurferCNN/utils/checkpoint.py | 24 +++++++++++----------- 5 files changed, 30 insertions(+), 26 deletions(-) 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/CorpusCallosum/shape/subsegment_contour.py b/CorpusCallosum/shape/subsegment_contour.py index 801489c00..2434081df 100644 --- a/CorpusCallosum/shape/subsegment_contour.py +++ b/CorpusCallosum/shape/subsegment_contour.py @@ -880,9 +880,8 @@ def transform_to_acpc_standard( norms_product = np.linalg.norm(ac_pc_vec) * np.linalg.norm(posterior_vector) theta = np.arccos(dot_product / norms_product) - # Determine the sign of the angle using the 2D cross product (z-component). - # NumPy 2.0 removed np.cross for 2-element vectors, so compute the scalar directly. - cross_product = ac_pc_vec[0] * posterior_vector[1] - ac_pc_vec[1] * posterior_vector[0] + # Determine the sign of the angle using cross product + cross_product = np.cross(ac_pc_vec, posterior_vector) if cross_product < 0: theta = -theta diff --git a/CorpusCallosum/utils/mapping_helpers.py b/CorpusCallosum/utils/mapping_helpers.py index b02dbee05..666e5e08b 100644 --- a/CorpusCallosum/utils/mapping_helpers.py +++ b/CorpusCallosum/utils/mapping_helpers.py @@ -84,9 +84,8 @@ def correct_nodding(ac_pt: Vector2d, pc_pt: Vector2d) -> AffineMatrix3x3: norms_product = np.linalg.norm(ac_pc_vec) * np.linalg.norm(posterior_vector) theta = np.arccos(dot_product / norms_product) - # Determine the sign of the angle using the 2D cross product (z-component). - # NumPy 2.0 removed np.cross for 2-element vectors, so compute the scalar directly. - cross_product = ac_pc_vec[0] * posterior_vector[1] - ac_pc_vec[1] * posterior_vector[0] + # Determine the sign of the angle using cross product + cross_product = np.cross(ac_pc_vec, posterior_vector) if cross_product < 0: theta = -theta diff --git a/FastSurferCNN/utils/checkpoint.py b/FastSurferCNN/utils/checkpoint.py index 39176eb74..701aaeaaa 100644 --- a/FastSurferCNN/utils/checkpoint.py +++ b/FastSurferCNN/utils/checkpoint.py @@ -15,11 +15,11 @@ # IMPORTS import os import sys -import tempfile from collections.abc import MutableSequence 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 @@ -401,20 +401,20 @@ def download_checkpoint( raise RuntimeError(message, responses) else: response = next(r for r in responses if r.ok) - # Write atomically: download to a temporary file in the same directory, then - # os.replace() it into place. This prevents a concurrent caller (e.g. the CC - # module loads its localization and segmentation models in parallel threads, - # both triggering download_checkpoints) from seeing a half-written file via the - # `checkpoint_path.exists()` guard and torch.load-ing a truncated checkpoint. checkpoint_path = Path(checkpoint_path) - fd, tmp_path = tempfile.mkstemp(dir=checkpoint_path.parent, suffix=".part") + temporary_path = checkpoint_path.with_name( + f".{checkpoint_path.name}.{uuid4().hex}.tmp" + ) try: - with os.fdopen(fd, "wb") as f: + # 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(tmp_path, checkpoint_path) - except BaseException: - Path(tmp_path).unlink(missing_ok=True) - raise + 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: