perf: optimize data / prediction hot paths

konfai/data/patching.py

@@ -176,12 +176,17 @@ def __init__(
         self.patch_size = patch_size
         self.patch_combine = patch_combine
         self.batch = batch
+        self._filled = 0
 
     def add_layer(self, index: int, layer: torch.Tensor) -> None:
+        if self._layer_accumulator[index] is None:
+            self._filled += 1
         self._layer_accumulator[index] = layer
 
     def is_full(self) -> bool:
-        return len(self.patch_slices) == len([v for v in self._layer_accumulator if v is not None])
+        # O(1): a running counter avoids re-scanning every slot after each added patch
+        # (the completion check ran once per patch, i.e. O(P^2) per case).
+        return self._filled == len(self.patch_slices)
 
     def assemble(self) -> torch.Tensor:
         n = 2 if self.batch else 1
@@ -195,7 +200,8 @@ def assemble(self) -> torch.Tensor:
         result = torch.zeros(
             (list(reference.shape[:n]) + list(max([[v.stop for v in patch] for patch in self.patch_slices]))),
             dtype=reference.dtype,
-        ).to(reference.device)
+            device=reference.device,
+        )
         # Overlap blending weights each patch (edge bands < 1 so interior overlaps sum to unity).
         # A voxel covered by fewer patches (a volume border without whole-image padding) would sum
         # to < 1 and come out darkened (x0.5 edges, x0.25 corners), so divide by the accumulated
@@ -224,6 +230,7 @@ def assemble(self) -> torch.Tensor:
         result = result[tuple([slice(None, None)] + [slice(0, s) for s in self.shape])]
 
         self._layer_accumulator.clear()
+        self._filled = 0
         return result
 
 

konfai/data/transform.py

@@ -163,8 +163,17 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute)
         min_value = float(min_value)
         max_value = float(max_value)
 
-        tensor[torch.where(tensor.float() < min_value)] = min_value
-        tensor[torch.where(tensor.float() > max_value)] = max_value
+        # Fast path: one fused in-place clamp instead of two float()-copy + where-scatter passes.
+        # Restricted to float32 (integer tensors reject float bounds; float16/float64 would compare
+        # at a different precision than the legacy float()-cast scatter) and to non-NaN bounds: a
+        # NaN bound — from a dynamic min/max/percentile over data containing NaN — makes clamp_
+        # propagate NaN to the whole tensor, whereas the legacy scatter no-ops on it (NaN
+        # comparisons are False). All other cases keep the exact original behaviour byte-for-byte.
+        if tensor.dtype == torch.float32 and min_value == min_value and max_value == max_value:
+            tensor.clamp_(min=min_value, max=max_value)
+        else:
+            tensor[torch.where(tensor.float() < min_value)] = min_value
+            tensor[torch.where(tensor.float() > max_value)] = max_value
         if self.save_clip_min:
             cache_attribute["Min"] = min_value
         if self.save_clip_max:

konfai/predictor.py

@@ -712,6 +712,10 @@ def __init__(self, model: Network, combine: Reduction):
         self._base_model_name = model.get_name()
         self._state_sources: list[dict[str, Any] | Path | str] = []
         self._loaded_state_index: int | None = None
+        # Cache the CPU state_dict per index so a local-path ensemble is read from
+        # disk once, not re-read + re-unpickled on every batch (the index cycles
+        # 0..N-1 each forward, so the next batch would otherwise reload all N).
+        self._state_cache: dict[int, dict[str, Any]] = {}
         self.add_module(
             self._model_name,
             copy.deepcopy(model),
@@ -732,10 +736,14 @@ def _read_state_source(self, source: dict[str, Any] | Path | str) -> dict[str, A
     def _ensure_model_loaded(self, index: int) -> Network:
         model = self._get_model()
         if self._loaded_state_index != index:
+            state = self._state_cache.get(index)
+            if state is None:
+                state = self._read_state_source(self._state_sources[index])
+                self._state_cache[index] = state
             # Checkpoints are keyed by the base model name, not by the streamed
             # ensemble suffix added after the previous load.
             model.set_name(self._base_model_name)
-            model.load(self._read_state_source(self._state_sources[index]), init=False)
+            model.load(state, init=False)
             model.set_name(f"{self._base_model_name}_{index}")
             self._loaded_state_index = index
         return model
@@ -749,6 +757,7 @@ def load(self, state_sources: list[dict[str, Any] | Path | str]):
         """
         self._state_sources = state_sources
         self._loaded_state_index = None
+        self._state_cache = {}
         if len(self._state_sources) == 1:
             self._ensure_model_loaded(0)
 

konfai/utils/dataset.py

@@ -976,7 +976,9 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) -
 
             path = self._path(name)
             info = get_dicom_info(path)
-            data, origin, spacing, direction = read_dicom_series_slice(path, slices, series_uid=info["series_uid"])
+            data, origin, spacing, direction = read_dicom_series_slice(
+                path, slices, series_uid=info["series_uid"], info=info
+            )
             info.update(origin=origin, spacing=spacing, direction=direction)
             return data, self._attributes(info)
 
@@ -986,13 +988,18 @@ def file_to_data_statistics(
             name: str,
             channels: list[int] | None = None,
         ) -> dict[str, float]:
-            shape, _ = self.get_infos(group, name)
+            from konfai.utils.dicom import get_dicom_info, read_dicom_series_slice
+
+            path = self._path(name)
+            info = get_dicom_info(path)
+            shape = info["shape"]
             state: dict[str, float] | None = None
             for index in range(shape[1]):
-                chunk, _ = self.file_to_data_slice(
-                    group,
-                    name,
+                chunk, _, _, _ = read_dicom_series_slice(
+                    path,
                     (slice(None), slice(index, index + 1), slice(None), slice(None)),
+                    series_uid=info["series_uid"],
+                    info=info,
                 )
                 if channels is not None:
                     chunk = chunk[channels]
@@ -1088,6 +1095,7 @@ def __init__(self, filename: str | Path, file_format: str) -> None:
         self.filename = str(filename)
         self.file_format = file_format
         self._names_cache: dict[str, list[str]] = {}
+        self._infos_cache: dict[tuple[str, str], tuple[list[int], Attribute]] = {}
 
     def _exists_on_disk(self) -> bool:
         if os.path.exists(self.filename):
@@ -1102,6 +1110,7 @@ def write(
         attributes: Attribute | None = None,
     ) -> None:
         self._names_cache.clear()
+        self._infos_cache.clear()
         if attributes is None:
             attributes = Attribute()
         if self.is_directory:
@@ -1290,6 +1299,15 @@ def get_group(self) -> list[str]:
         return list(groups)
 
     def get_infos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
+        # Memoize the header read (SITK reader + ReadImageInformation, or the HDF5/Zarr
+        # metadata parse): get_infos is called once per name per group per build-pass at
+        # setup, so caching it (like get_names) avoids re-parsing the same header N times.
+        # Cache and hand back copies so a caller mutating the geometry cannot poison it.
+        cache_key = (groups, name)
+        cached = self._infos_cache.get(cache_key)
+        if cached is not None:
+            shape, attr = cached
+            return list(shape), Attribute(attr)
         if self.is_directory:
             for sub_directory in self._get_sub_directories(groups):
                 group = groups.split("/")[-1]
@@ -1300,10 +1318,14 @@ def get_infos(self, groups: str, name: str) -> tuple[list[int], Attribute]:
                         self.file_format,
                         self.level,
                     ) as file:
-                        return file.get_infos("", group)
+                        result = file.get_infos("", group)
+                        self._infos_cache[cache_key] = (list(result[0]), Attribute(result[1]))
+                        return result
         else:
             with Dataset.File(self.filename, True, self.file_format, self.level) as file:
-                return file.get_infos(groups, name)
+                result = file.get_infos(groups, name)
+                self._infos_cache[cache_key] = (list(result[0]), Attribute(result[1]))
+                return result
         raise NameError(f"Dataset entry '{groups}/{name}' not found in {self.filename}.")
 
     def get_statistics(self, groups: str) -> dict[str, dict[str, dict[str, float | list[float]]]]:

konfai/utils/dicom.py

@@ -387,6 +387,7 @@ def get_dicom_info(
     return {
         "series_uid": selected_uid,
         "files": files,
+        "sorted_files": [Path(ds.filename) for ds in datasets],
         "shape": [1, len(datasets), rows, columns],
         "origin": origin,
         "spacing": spacing,
@@ -400,9 +401,15 @@ def read_dicom_series_slice(
     *,
     series_uid: str | None = None,
     apply_rescale: bool = True,
+    info: dict[str, Any] | None = None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Read only the selected DICOM slices and return updated patch geometry."""
-    info = get_dicom_info(directory, series_uid=series_uid)
+    if info is None:
+        info = get_dicom_info(directory, series_uid=series_uid)
+    elif series_uid is not None and series_uid != info["series_uid"]:
+        raise DatasetManagerError(
+            f"series_uid '{series_uid}' does not match the provided info series '{info['series_uid']}'."
+        )
     shape = info["shape"]
     if len(slices) != len(shape):
         raise DatasetManagerError(f"Expected {len(shape)} slices, got {len(slices)}.")
@@ -411,10 +418,8 @@ def read_dicom_series_slice(
     if list(channel_indices) not in ([0], []):
         raise DatasetManagerError("DICOM stores scalar data and supports only channel 0.")
 
-    _selected_uid, files = _select_series_files(directory, series_uid or info["series_uid"])
-    headers = sort_series(files, stop_before_pixels=True)
     z_indices = list(range(*normalized[1].indices(shape[1])))
-    selected_files = [Path(headers[index].filename) for index in z_indices]
+    selected_files = [info["sorted_files"][index] for index in z_indices]
     datasets = sort_series(selected_files)
     volume = read_volume(datasets, apply_rescale=apply_rescale)
     volume = volume[normalized[0], :, normalized[2], normalized[3]]