Coverage for src/crackmnist_eccv/utility.py: 53%

1# SPDX-FileCopyrightText: 2026 German Aerospace Center (DLR) 

2# SPDX-License-Identifier: GPL-3.0-or-later 

3 

4"""Pollution utility functions for CrackMNIST experiments. 

5 

6Provides per-image noise injection (Gaussian, Salt & Pepper, Gaussian Blur, 

7Blob/MVS, Uniform) and crack-tip mask enlargement used during training and 

8evaluation. 

9""" 

10 

11from __future__ import annotations 

12 

13import math 

14import random 

15 

16import numpy as np 

17from scipy.interpolate import griddata 

18import torch 

19import torch.nn.functional as F 

20 

21 

22def add_gaussian_noise( 

23 tensor: torch.Tensor, mean: float = 0.0, intensity: float = 0.5 

24) -> torch.Tensor: 

25 """Add Gaussian noise to a tensor. 

26 

27 Args: 

28 tensor: Input tensor. 

29 mean: Mean of the Gaussian noise. 

30 intensity: Standard deviation of the noise. 

31 

32 Returns: 

33 Noisy tensor of the same shape. 

34 """ 

35 noise = torch.randn_like(tensor) * intensity + mean 

36 return tensor + noise 

37 

38 

39def add_salt_and_pepper_noise( 

40 tensor: torch.Tensor, 

41 intensity: float = 0.25, 

42 salt_to_pepper_ratio: float = 0.5, 

43) -> torch.Tensor: 

44 """Add salt-and-pepper noise to a tensor. 

45 

46 Args: 

47 tensor: Input tensor. 

48 intensity: Fraction of pixels to corrupt. 

49 salt_to_pepper_ratio: Proportion of corrupted pixels set to 1.0 

50 (salt) vs 0.0 (pepper). 

51 

52 Returns: 

53 Noisy tensor of the same shape. 

54 """ 

55 noisy_tensor = tensor.clone() 

56 total_pixels = tensor.numel() 

57 num_noisy_pixels = int(intensity * total_pixels) 

58 

59 num_salt = int(salt_to_pepper_ratio * num_noisy_pixels) 

60 num_pepper = num_noisy_pixels - num_salt 

61 

62 salt_indices = torch.randint(0, total_pixels, (num_salt,)) 

63 pepper_indices = torch.randint(0, total_pixels, (num_pepper,)) 

64 

65 noisy_tensor.view(-1)[salt_indices] = 1.0 

66 noisy_tensor.view(-1)[pepper_indices] = 0.0 

67 

68 return noisy_tensor 

69 

70 

71def add_uniform_noise(tensor: torch.Tensor, intensity: float = 0.1) -> torch.Tensor: 

72 """Add zero-centred uniform noise in [-intensity, +intensity). 

73 

74 Args: 

75 tensor: Input tensor. 

76 intensity: Half-range of the uniform distribution. Must be >= 0. 

77 

78 Returns: 

79 Noisy tensor of the same shape. 

80 

81 Raises: 

82 ValueError: If intensity is negative. 

83 """ 

84 if intensity < 0: 

85 raise ValueError("intensity must be non-negative") 

86 

87 noise = torch.empty_like(tensor).uniform_(-intensity, intensity) 

88 return tensor + noise 

89 

90 

91def apply_gaussian_blur(tensor: torch.Tensor, intensity: float = 4.0) -> torch.Tensor: 

92 """Apply channel-wise Gaussian blur to a (B, C, H, W) tensor. 

93 

94 Kernel size is derived from sigma following Yoshihara et al. (2023): 

95 ``kernel_size = round(8 * intensity + 1)``. 

96 

97 Args: 

98 tensor: Input tensor of shape (B, C, H, W). 

99 intensity: Standard deviation (sigma) of the Gaussian kernel. 

100 

101 Returns: 

102 Blurred tensor of the same shape. 

103 """ 

104 kernel_size = round(8 * intensity + 1) 

105 

106 if kernel_size % 2 == 0: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 kernel_size += 1 

108 

109 def gaussian_kernel(kernel_size: int, intensity: float) -> torch.Tensor: 

110 """Build a normalised 2-D Gaussian kernel. 

111 

112 Args: 

113 kernel_size: Size of the square kernel (must be odd). 

114 intensity: Standard deviation (sigma) of the Gaussian. 

115 

116 Returns: 

117 Normalised kernel tensor of shape (1, 1, kernel_size, kernel_size). 

118 """ 

119 kernel = torch.arange(kernel_size).float() - (kernel_size - 1) / 2 

120 x_grid, y_grid = torch.meshgrid(kernel, kernel, indexing="ij") 

121 kernel_2d = torch.exp(-(x_grid**2 + y_grid**2) / (2 * intensity**2)) 

122 kernel_2d = kernel_2d / kernel_2d.sum() 

123 return kernel_2d.unsqueeze(0).unsqueeze(0) 

124 

125 kernel = gaussian_kernel(kernel_size, intensity).to(tensor.device) 

126 

127 blurred_tensor = tensor.clone() 

128 for channel in range(tensor.shape[1]): 

129 tensor_channel = tensor[:, channel : channel + 1, :, :] 

130 blurred_tensor_channel = F.conv2d( 

131 tensor_channel, kernel, padding=kernel_size // 2 

132 ) 

133 blurred_tensor[:, channel : channel + 1, :, :] = blurred_tensor_channel 

134 

135 return blurred_tensor 

136 

137 

138# BLOB POLLUTION SEGMENT (MVS) ### 

139def generate_blob_mask_to_match_oar( 

140 shape: tuple[int, int], 

141 *, 

142 radius: int = 4, 

143 target_min_oar: float = 0.10, 

144 target_max_oar: float = 0.30, 

145 max_draw_attempts: int = 2_000, 

146 rng: np.random.Generator | None = None, 

147) -> tuple[np.ndarray, float]: 

148 """Create a binary mask whose OAR lies in [target_min_oar, target_max_oar]. 

149 

150 Keeps stamping discs until the lower bound is reached, skipping discs 

151 that would push the OAR past ``target_max_oar``. 

152 

153 Args: 

154 shape: (H, W) resolution of the mask. 

155 radius: Intended radius of each disc in pixels. 

156 target_min_oar: Lower bound of the desired OAR interval. 

157 target_max_oar: Upper bound of the desired OAR interval. 

158 max_draw_attempts: Safety cap on total centre-samples tried. 

159 rng: Optional RNG for deterministic tests. 

160 

161 Returns: 

162 Tuple of (mask, oar) where mask is a uint8 array and oar is 

163 ``mask.mean()``. 

164 

165 Raises: 

166 ValueError: If target_min_oar > target_max_oar. 

167 RuntimeError: If max_draw_attempts is exceeded. 

168 """ 

169 if target_min_oar > target_max_oar: 

170 raise ValueError(f"Invalid OAR range: {target_min_oar} > {target_max_oar}") 

171 

172 rng = rng or np.random.default_rng() 

173 H, W = shape 

174 total_px = H * W 

175 

176 mask = np.zeros((H, W), dtype=np.uint8) 

177 current_px = 0 

178 

179 Y, X = np.ogrid[:H, :W] 

180 

181 for _ in range(max_draw_attempts): 181 ↛ 208line 181 didn't jump to line 208 because the loop on line 181 didn't complete

182 current_oar = current_px / total_px 

183 if current_oar >= target_min_oar: 

184 break 

185 

186 remaining_px_max = target_max_oar * total_px - current_px 

187 if remaining_px_max < 1: 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true

188 break 

189 

190 r_use = min( 

191 radius, 

192 math.floor(math.sqrt(remaining_px_max / math.pi)) or 1, 

193 ) 

194 

195 cx = rng.integers(0, W) 

196 cy = rng.integers(0, H) 

197 

198 disc = (X - cx) ** 2 + (Y - cy) ** 2 <= r_use**2 

199 new_pixels = disc & (mask == 0) 

200 

201 added = int(new_pixels.sum()) 

202 if added == 0: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true

203 continue 

204 

205 mask[new_pixels] = 1 

206 current_px += added 

207 else: 

208 raise RuntimeError( 

209 "Exceeded max_draw_attempts while trying to reach the OAR band " 

210 f"[{target_min_oar}, {target_max_oar}]. " 

211 "Consider increasing `max_draw_attempts`, enlarging `radius`, " 

212 "or relaxing the OAR constraints." 

213 ) 

214 

215 return mask, mask.mean() 

216 

217 

218def add_blob_pollution_and_interp( 

219 tensor: torch.Tensor, 

220 *, 

221 radius: int = 4, 

222 target_min_oar: float = 0.10, 

223 target_max_oar: float = 0.20, 

224 interp_method: str = "linear", 

225) -> tuple[torch.Tensor, np.ndarray, float]: 

226 """Mask a (1, C, H, W) tensor with blob occlusion and in-paint the holes. 

227 

228 Args: 

229 tensor: Input tensor of shape (1, C, H, W). 

230 radius: Disc radius for blob generation. 

231 target_min_oar: Lower OAR bound. 

232 target_max_oar: Upper OAR bound. 

233 interp_method: Interpolation method passed to scipy.griddata. 

234 

235 Returns: 

236 Tuple of (tensor_filled, blob_mask, oar) where tensor_filled has NaNs 

237 replaced by interpolated values, blob_mask is a (H, W) uint8 array, 

238 and oar is the final Occluded Area Ratio. 

239 

240 Raises: 

241 ValueError: If tensor shape is not (1, C, H, W). 

242 RuntimeError: If fewer than 4 valid points remain for interpolation. 

243 """ 

244 if tensor.ndim != 4 or tensor.size(0) != 1: 

245 raise ValueError("Expected tensor shape (1, C, H, W).") 

246 

247 dev = tensor.device 

248 _, C, H, W = tensor.shape 

249 

250 blob_mask, oar = generate_blob_mask_to_match_oar( 

251 (H, W), 

252 radius=radius, 

253 target_min_oar=target_min_oar, 

254 target_max_oar=target_max_oar, 

255 ) 

256 blob_mask_t = torch.from_numpy(blob_mask).to(dev) 

257 

258 polluted = tensor.clone() 

259 nan_val = torch.tensor(float("nan"), device=dev) 

260 polluted[:, :, blob_mask_t.bool()] = nan_val 

261 

262 yy, xx = np.mgrid[0:H, 0:W] 

263 xx_flat, yy_flat = xx.ravel(), yy.ravel() 

264 

265 filled = polluted.clone() 

266 for c in range(C): 

267 img = polluted[0, c].cpu().numpy() 

268 valid = ~np.isnan(img) 

269 n_valid = int(valid.sum()) 

270 

271 if n_valid < 4: 

272 raise RuntimeError( 

273 f"Channel {c}: only {n_valid} valid points left; cannot interpolate." 

274 ) 

275 

276 pts = np.column_stack((xx_flat[valid.ravel()], yy_flat[valid.ravel()])) 

277 vals = img[valid] 

278 

279 method = interp_method if n_valid >= 30 else "nearest" 

280 interp = griddata(pts, vals, (xx, yy), method=method) 

281 

282 if np.isnan(interp).any(): 

283 nearest = griddata(pts, vals, (xx, yy), method="nearest") 

284 interp = np.where(np.isnan(interp), nearest, interp) 

285 

286 filled[0, c] = torch.from_numpy(interp).to(dev) 

287 

288 return filled, blob_mask, oar 

289 

290 

291def apply_per_image_pollution( 

292 inputs: torch.Tensor, 

293 pollution_ratio: float = 0.5, 

294 pollution_type: str = "gaussian", 

295 **kwargs: float | str, 

296) -> torch.Tensor: 

297 """Apply noise to a random subset of images in a batch. 

298 

299 Args: 

300 inputs: Batch of input tensors (B, C, H, W). 

301 pollution_ratio: Proportion of images to pollute. 

302 pollution_type: One of 'gaussian', 'salt_pepper', 'gaussian_blur', 

303 'blob', 'uniform'. 

304 kwargs: Extra keyword arguments forwarded to the pollution function. 

305 Values must be float, int, or str. 

306 

307 Returns: 

308 Tensor with selected images polluted, same shape as inputs. 

309 

310 Raises: 

311 ValueError: If pollution_type is not recognised. 

312 """ 

313 polluted_inputs = inputs.clone() 

314 batch_size = inputs.size(0) 

315 num_polluted = int(batch_size * pollution_ratio) 

316 

317 indices_to_pollute = sorted(random.sample(range(batch_size), num_polluted)) 

318 

319 for i in indices_to_pollute: 

320 img = inputs[i].unsqueeze(0) 

321 

322 if pollution_type == "gaussian": 

323 img = add_gaussian_noise(img, **kwargs) # ty: ignore[invalid-argument-type] 

324 elif pollution_type == "salt_pepper": 

325 img = add_salt_and_pepper_noise(img, **kwargs) # ty: ignore[invalid-argument-type] 

326 elif pollution_type == "gaussian_blur": 

327 img = apply_gaussian_blur(img, **kwargs) # ty: ignore[invalid-argument-type] 

328 elif pollution_type == "blob": 

329 img, _blob_mask, _oar = add_blob_pollution_and_interp(img, **kwargs) # ty: ignore[invalid-argument-type] 

330 elif pollution_type == "uniform": 

331 img = add_uniform_noise(img, **kwargs) # ty: ignore[invalid-argument-type] 

332 else: 

333 raise ValueError(f"Unknown pollution type: {pollution_type}") 

334 

335 polluted_inputs[i] = img.squeeze(0) 

336 

337 return polluted_inputs 

338 

339 

340def enlarge_crack_tip(mask: np.ndarray, kernel_size: int = 3) -> np.ndarray: 

341 """Enlarge the annotated crack tip region in a binary mask. 

342 

343 Expands each foreground pixel into a square neighbourhood of size 

344 ``kernel_size x kernel_size``. 

345 

346 Args: 

347 mask: 2D numpy array of shape (H, W) with foreground pixels encoded 

348 as 255 and background as 0. 

349 kernel_size: Size of the square neighbourhood. Must be an odd integer 

350 (e.g. 1, 3, 5). Default is 3. 

351 

352 Returns: 

353 New 2D mask of the same shape with the crack tip region enlarged. 

354 """ 

355 pad = kernel_size // 2 

356 height, width = mask.shape 

357 new_mask = np.copy(mask) 

358 y, x = np.where(mask == 255) 

359 for i in range(len(x)): 

360 y_start = max(y[i] - pad, 0) 

361 y_end = min(y[i] + pad + 1, height) 

362 x_start = max(x[i] - pad, 0) 

363 x_end = min(x[i] + pad + 1, width) 

364 new_mask[y_start:y_end, x_start:x_end] = 255 

365 return new_mask