Coverage for supercat/utils.py: 100.00%
93 statements
« prev ^ index » next coverage.py v7.10.6, created at 2026-09-11 13:26 +0000
1from pathlib import Path
2import torch
5def generate_intervals(n: int, k: int) -> list[tuple[int, int]]:
6 """
7 Generate intervals of size `k` that fit within a range of size `n`.
9 Args:
10 n (int): The size of the range.
11 k (int): The size of each interval.
13 Returns:
14 list[tuple[int, int]]: A list of tuples representing the intervals. Each tuple contains the start and end of an interval.
15 If `k` is greater than `n`, returns an empty list.
17 Raises:
18 ValueError: If `k` is larger than `n`, a ValueError is raised.
20 Examples:
21 >>> generate_intervals(10, 2)
22 [(0, 2), (2, 4), (4, 6), (6, 8), (8, 10)]
24 >>> generate_intervals(10, 3)
25 [(0, 3), (3, 6), (6, 9)]
27 >>> generate_intervals(5, 6)
28 []
29 """
30 if k > n or k == 0 or n == 0:
31 return []
33 # Calculate number of full intervals that fit into n
34 num_intervals = n // k
36 # Calculate remaining space (gap) after placing intervals
37 total_gap = n - (num_intervals * k)
39 # Distribute the gap equally at the start and end, with any remainder in the middle
40 start_gap = total_gap // 2
42 intervals = [
43 (start_gap + i * k, start_gap + (i + 1) * k) for i in range(num_intervals)
44 ]
46 return intervals
49def generate_overlapping_intervals(
50 total: int,
51 interval_size: int,
52 min_overlap: int,
53 check: bool = True,
54 variable_size: bool = False,
55) -> list[tuple[int, int]]:
56 """
57 Creates a list of overlapping intervals within a specified range, adjusting the interval size to ensure
58 that the overlap is approximately the same across all intervals.
60 Args:
61 total (int): The total range within which intervals are to be created.
62 max_interval_size (int): The maximum size of each interval.
63 min_overlap (int): The minimum number of units by which consecutive intervals overlap.
64 check (bool): If True, checks are performed to ensure that the intervals meet the specified conditions.
66 Returns:
67 list[tuple[int, int]]: A list of tuples where each tuple represents the start (inclusive)
68 and end (exclusive) of an interval.
70 Example:
71 >>> generate_overlapping_intervals(20, 5, 2)
72 [(0, 5), (3, 8), (6, 11), (9, 14), (12, 17), (15, 20)]
73 """
74 intervals = []
75 start = 0
77 if total == 0:
78 return intervals
80 max_interval_size = interval_size
81 assert interval_size
82 assert min_overlap is not None
83 assert (
84 interval_size > min_overlap
85 ), f"Max interval size of {interval_size} must be greater than min overlap of {min_overlap}"
87 # Calculate the number of intervals needed to cover the range
88 num_intervals, remainder = divmod(total - min_overlap, interval_size - min_overlap)
89 if remainder > 0:
90 num_intervals += 1
92 # Calculate the exact interval size to ensure consistent overlap
93 overlap = min_overlap
94 if variable_size:
95 if num_intervals > 1:
96 interval_size, remainder = divmod(
97 total + (num_intervals - 1) * overlap, num_intervals
98 )
99 if remainder > 0:
100 interval_size += 1
101 else:
102 # If the size is fixed, then vary the overlap to keep it even.
103 # num_intervals is the ceiling of (total - min_overlap)/(interval_size - min_overlap),
104 # so num_intervals * interval_size - total >= min_overlap * (num_intervals - 1)
105 # and the overlap below can never fall under min_overlap.
106 if num_intervals > 1:
107 overlap, remainder = divmod(
108 num_intervals * interval_size - total, num_intervals - 1
109 )
111 while True:
112 end = start + interval_size
113 if end > total:
114 end = total
115 start = max(end - interval_size, 0)
116 intervals.append((start, end))
117 start += interval_size - overlap
118 if end >= total:
119 break
121 if check:
122 assert intervals[0][0] == 0
123 assert intervals[-1][1] == total
124 assert (
125 len(intervals) == num_intervals
126 ), f"Expected {num_intervals} intervals, got {len(intervals)}"
128 assert (
129 interval_size <= max_interval_size
130 ), f"Interval size of {interval_size} exceeds max interval size of {max_interval_size}"
131 for interval in intervals:
132 assert (
133 interval[1] - interval[0] == interval_size
134 ), f"Interval size of {interval[1] - interval[0]} is not the expected size {interval_size}"
136 for i in range(1, len(intervals)):
137 overlap = intervals[i - 1][1] - intervals[i][0]
138 assert (
139 overlap >= min_overlap
140 ), f"Min overlap condition of {min_overlap} not met for intervals {intervals[i - 1]} and {intervals[i]} (overlap {overlap})"
142 return intervals
145def distance_to_boundary(size_i: int, size_j: int, size_k: int) -> torch.Tensor:
146 """
147 Builds a weight tensor where each element represents the minimum distance
148 to any edge of a 3D volume of size (size_i, size_j, size_k).
150 Args:
151 size_i (int): The size of the first dimension.
152 size_j (int): The size of the second dimension.
153 size_k (int): The size of the third dimension.
155 Returns:
156 torch.Tensor: A tensor of shape (size_i, size_j, size_k) where each
157 element is the minimum distance to any edge.
158 """
159 x = torch.arange(size_i).view(-1, 1, 1).expand(size_i, size_j, size_k)
160 y = torch.arange(size_j).view(1, -1, 1).expand(size_i, size_j, size_k)
161 z = torch.arange(size_k).view(1, 1, -1).expand(size_i, size_j, size_k)
163 # Compute the distance to the closest edge for each coordinate
164 distance_x = torch.minimum(x, size_i - 1 - x)
165 distance_y = torch.minimum(y, size_j - 1 - y)
166 distance_z = torch.minimum(z, size_k - 1 - z)
168 # Combine the distances to get the minimum distance to any edge
169 return torch.minimum(torch.minimum(distance_x, distance_y), distance_z)
172def write_image(data: torch.Tensor, path: Path | str) -> None:
173 """
174 Saves a 2D image or 3D volume data to the specified path in either .pt or image format.
176 Args:
177 data (torch.Tensor): The 3D tensor data to be saved.
178 path (Path | str): The destination file path where the data will be saved. If the suffix is
179 '.pt', the data will be saved as a PyTorch tensor. Otherwise, it will be saved as an image file.
181 Returns:
182 None
184 Notes:
185 - If the directory specified in the path does not exist, it will be created.
186 - Supports saving in '.pt' format for PyTorch tensors. Other formats will
187 be saved using `io.imsave` from skimage.io after converting the tensor to a NumPy array.
188 """
189 path = Path(path)
190 print(f"Writing to {path}")
191 path.parent.mkdir(exist_ok=True, parents=True)
192 suffix = path.suffix.lower()
193 match suffix:
194 case ".nrrd":
195 import nrrd
196 import nrrd.writer
198 # Add new type mappings
199 nrrd.writer._TYPEMAP_NUMPY2NRRD.update(
200 {
201 "f2": "float16",
202 }
203 )
204 nrrd.write(str(path), data.numpy())
205 case ".pt":
206 torch.save(data, path)
207 case ".tif":
208 import numpy as np
209 import tifffile as tiff
211 if isinstance(data, torch.Tensor):
212 data = data.numpy()
214 # Clip data to valid range
215 data = np.clip(data, -1.0, 1.0)
217 # Save output
218 data = (data + 1.0) * 255.0 / 2.0
219 data = data.astype(np.uint8)
220 tiff.imwrite(str(path), data, dtype=np.uint8)
221 case _:
222 import numpy as np
224 if isinstance(data, torch.Tensor):
225 data = data.numpy()
227 # Clip data to valid range
228 data = np.clip(data, -1.0, 1.0)
230 # Save output
231 data = (data + 1.0) * 255.0 / 2.0
232 data = data.astype(np.uint8)
234 if suffix == ".mat":
235 import hdf5storage
237 DEEPROCK_HDF5_KEY = "temp"
239 hdf5storage.savemat(
240 str(path),
241 {DEEPROCK_HDF5_KEY: data},
242 format="7.3",
243 oned_as="column",
244 store_python_metadata=True,
245 )
246 else:
247 from skimage import io
249 io.imsave(path, data)