|
10 | 10 | from alive_progress import alive_bar |
11 | 11 | from scipy.ndimage import zoom |
12 | 12 |
|
13 | | - |
14 | 13 | class DataReader: |
15 | 14 | def __init__(self, path: str | Path): |
16 | 15 | """ |
@@ -109,81 +108,65 @@ def load_volume( |
109 | 108 | unbinned_shape: tuple[int, int, int] | None = None, |
110 | 109 | ) -> np.ndarray: |
111 | 110 | """ |
112 | | - Loads the volume data based on the detected volume type. |
| 111 | + Loads the volume and resizes it to unbinned_shape if provided, using fast scipy.ndimage.zoom. |
113 | 112 |
|
114 | 113 | Args: |
115 | 114 | start_index (int): Start index for slicing (for stacks). |
116 | 115 | end_index (int): End index for slicing (for stacks). If None, loads the entire stack. |
117 | | - unbinned_shape (tuple): Shape of the volume without downsampling. Default is None (no binning). |
| 116 | + unbinned_shape (tuple): Desired shape (Z, Y, X). If None, no resizing is done. |
118 | 117 |
|
119 | 118 | Returns: |
120 | | - np.ndarray: Loaded (and possibly resized) volume data. |
| 119 | + np.ndarray: Loaded volume. |
121 | 120 | """ |
122 | | - |
123 | 121 | if end_index is None: |
124 | 122 | end_index = self.shape[0] |
125 | 123 |
|
126 | | - # compute how much we’ll need to zoom later |
127 | | - binning_factor = 1.0 |
128 | | - if unbinned_shape is not None: |
129 | | - binning_factor = unbinned_shape[0] / self.shape[0] |
130 | | - print(f"Mask bining factor: {binning_factor}") |
| 124 | + # Decide if resize is needed |
| 125 | + need_resize = False |
| 126 | + if unbinned_shape is not None and self.shape != unbinned_shape: |
| 127 | + need_resize = True |
| 128 | + zoom_factors = tuple(u / s for u, s in zip(unbinned_shape, self.shape)) |
| 129 | + print(f"Zoom factors: {zoom_factors}") |
| 130 | + else: |
| 131 | + zoom_factors = (1.0, 1.0, 1.0) |
131 | 132 |
|
132 | | - # if we’re going to zoom, adjust the requested slice range to include padding |
133 | | - if binning_factor != 1.0: |
| 133 | + # Optional padding if resizing |
| 134 | + if need_resize: |
134 | 135 | start_index_ini, end_index_ini = start_index, end_index |
135 | | - start_index = int(start_index_ini / binning_factor) - 1 |
| 136 | + start_index = int(start_index_ini / zoom_factors[0]) - 1 |
136 | 137 | start_index = max(start_index, 0) |
137 | | - end_index = int(end_index_ini / binning_factor) + 1 |
| 138 | + end_index = int(end_index_ini / zoom_factors[0]) + 1 |
138 | 139 | end_index = min(end_index, self.shape[0]) |
139 | | - print( |
140 | | - f"Mask start index padded: {start_index} - Mask end index padded: {end_index}" |
141 | | - ) |
| 140 | + print(f"Volume start index padded: {start_index} - end: {end_index}") |
142 | 141 |
|
143 | | - # load the raw volume |
| 142 | + # Load volume from stack or mhd |
144 | 143 | if not self.volume_info["stack"]: |
145 | | - # single-file (e.g. .mhd) |
146 | 144 | if self.volume_info["type"] == "mhd": |
147 | 145 | volume, _ = _load_raw_data_with_mhd(self.path) |
148 | 146 | volume = volume[start_index:end_index, :, :] |
149 | 147 | else: |
150 | | - # image stack |
151 | 148 | volume = self._load_image_stack( |
152 | 149 | self.volume_info["file_list"], start_index, end_index |
153 | 150 | ) |
154 | 151 |
|
155 | | - # if we need to upsample + crop + resize each slice back to unbinned_shape: |
156 | | - if binning_factor != 1.0 and unbinned_shape is not None: |
157 | | - print("Resizing mask") |
158 | | - # 1) zoom the 3D block |
159 | | - volume = zoom(volume, zoom=binning_factor, order=0) |
160 | | - |
161 | | - # 2) figure out where our original slice window ended up |
162 | | - start_up = int(abs(start_index * binning_factor - start_index_ini)) |
163 | | - end_up = start_up + (end_index_ini - start_index_ini) |
164 | | - start_up = max(start_up, 0) |
165 | | - end_up = min(end_up, volume.shape[0]) |
166 | | - |
167 | | - # 3) crop to just those slices |
168 | | - volume = volume[start_up:end_up, :, :] |
169 | | - |
170 | | - # 4) allocate exactly the target (unbinned) shape and resize each slice |
171 | | - volume_resized = np.empty( |
172 | | - (volume.shape[0], unbinned_shape[1], unbinned_shape[2]), |
173 | | - dtype=volume.dtype, |
174 | | - ) |
175 | | - for i in range(volume.shape[0]): |
176 | | - volume_resized[i] = cv2.resize( |
177 | | - volume[i], |
178 | | - (unbinned_shape[2], unbinned_shape[1]), |
179 | | - interpolation=cv2.INTER_LINEAR, |
180 | | - ) |
181 | | - |
182 | | - # 5) replace volume with the resized version |
183 | | - volume = volume_resized |
184 | | - |
| 152 | + if need_resize: |
| 153 | + print("Resizing with scipy.ndimage.zoom...") |
| 154 | + |
| 155 | + # Ensure float32 for better memory and speed |
| 156 | + volume = volume.astype(np.float32) |
| 157 | + volume = zoom(volume, zoom=zoom_factors, order=0) # Nearest-neighbor for mask |
| 158 | + |
| 159 | + # Final crop to original range |
| 160 | + crop_start = int(abs(start_index * zoom_factors[0] - start_index_ini)) |
| 161 | + crop_end = crop_start + (end_index_ini - start_index_ini) |
| 162 | + crop_start = max(crop_start, 0) |
| 163 | + crop_end = min(crop_end, volume.shape[0]) |
| 164 | + |
| 165 | + volume = volume[crop_start:crop_end, :, :] |
| 166 | + |
185 | 167 | return volume |
186 | 168 |
|
| 169 | + |
187 | 170 | def _custom_image_reader(self, file_path: Path) -> np.ndarray: |
188 | 171 | """ |
189 | 172 | Reads an image from the given file path. |
@@ -380,3 +363,4 @@ def _load_raw_data_with_mhd( |
380 | 363 | # End 3D fix |
381 | 364 |
|
382 | 365 | return (data, meta_dict) |
| 366 | + |
0 commit comments