Coverage for supercat/apps.py: 100.00%

1from pathlib import Path 

2from widitapp import WiDiTApp 

3import cluey 

4from cluey import main, method, tool 

5from contextlib import nullcontext 

6from rich.progress import Progress 

7 

8from .data import build_datasets3D, build_datasets2D 

9 

10 

11class Supercat(WiDiTApp): 

12 @method 

13 def datasets( 

14 self, 

15 dim: int = cluey.Option(3, help="Number of spatial dimensions (2 or 3)"), 

16 deeprock: Path = cluey.Option( 

17 None, help="Path to the DeepRock dataset directory" 

18 ), 

19 scale: int = cluey.Option(4, help="Scale factor for downsampling the images"), 

20 augment: bool = cluey.Option( 

21 True, help="Apply data augmentation to the training images" 

22 ), 

23 **kwargs, 

24 ) -> tuple: 

25 """Build training and validation datasets for 2D or 3D super-resolution.""" 

26 build_function = build_datasets2D if dim == 2 else build_datasets3D 

27 return build_function(deeprock=deeprock, scale=scale, train_augment=augment) 

28 

29 @main 

30 def predict( 

31 self, 

32 input: Path = cluey.Option(None, help="Path to the input image"), 

33 output: Path = cluey.Option(None, help="Path to save the predicted image"), 

34 size: int = cluey.Option( 

35 100, help="Default prediction tile size for all spatial dimensions" 

36 ), 

37 size_i: int = cluey.Option( 

38 0, help="Prediction tile size along the i dimension (0 uses size)" 

39 ), 

40 size_j: int = cluey.Option( 

41 0, help="Prediction tile size along the j dimension (0 uses size)" 

42 ), 

43 size_k: int = cluey.Option( 

44 0, help="Prediction tile size along the k dimension (0 uses size)" 

45 ), 

46 overlap: int = cluey.Option( 

47 10, help="Default overlap between prediction tiles in pixels or voxels" 

48 ), 

49 overlap_i: int = cluey.Option( 

50 0, help="Tile overlap along the i dimension (0 uses overlap)" 

51 ), 

52 overlap_j: int = cluey.Option( 

53 0, help="Tile overlap along the j dimension (0 uses overlap)" 

54 ), 

55 overlap_k: int = cluey.Option( 

56 0, help="Tile overlap along the k dimension (0 uses overlap)" 

57 ), 

58 checkpoint: Path = cluey.Option(None, help="Path to the model checkpoint"), 

59 num_sampling_steps: int = cluey.Option( 

60 250, help="Number of diffusion sampling steps per prediction" 

61 ), 

62 seed: int = cluey.Option(42, help="Random seed for diffusion sampling"), 

63 single_crop: bool = cluey.Option(False, help="Predict only the center tile"), 

64 **kwargs, 

65 ): 

66 """Generate and save a super-resolution prediction for an input image.""" 

67 import torch 

68 from widit import load_model 

69 

70 from .data import read_image_as_tensor 

71 from .models import DiffusionPredictionModel 

72 from .utils import write_image 

73 

74 torch.set_grad_enabled(False) 

75 device = "cuda" if torch.cuda.is_available() else "cpu" 

76 

77 assert output is not None, "Must provide output path" 

78 output_path = Path(output) 

79 output_path.parent.mkdir(parents=True, exist_ok=True) 

80 

81 assert checkpoint is not None, "Must provide path to model checkpoint" 

82 

83 size_i = size_i or size 

84 size_j = size_j or size 

85 size_k = size_k or size 

86 

87 overlap_i = overlap_i or overlap 

88 overlap_j = overlap_j or overlap 

89 overlap_k = overlap_k or overlap 

90 

91 assert input is not None, "Must provide input path" 

92 input_image = read_image_as_tensor(input, size=(size_k, size_j, size_i)) 

93 spatial_dims = input_image.ndim - 1 

94 assert spatial_dims in ( 

95 2, 

96 3, 

97 ), f"Input image must have 3 or 4 dimensions (C, [D], H, W), got {input_image.shape}" 

98 

99 assert ( 

100 input_image.shape[-1] == size_i 

101 ), f"Input image size in i dimension ({input_image.shape[-1]}) does not match specified size_i ({size_i})" 

102 assert ( 

103 input_image.shape[-2] == size_j 

104 ), f"Input image size in j dimension ({input_image.shape[-2]}) does not match specified size_j ({size_j})" 

105 if spatial_dims == 3: 

106 assert ( 

107 input_image.shape[-3] == size_k 

108 ), f"Input image size in k dimension ({input_image.shape[-3]}) does not match specified size_k ({size_k})" 

109 

110 # Load checkpoint 

111 model = load_model(checkpoint) 

112 

113 diffusion = model.out_channels == 2 

114 if diffusion: 

115 model = DiffusionPredictionModel(model, num_sampling_steps) 

116 torch.manual_seed(seed) 

117 

118 model.to(device=device) 

119 model.eval() 

120 

121 prediction = self.generate_prediction( 

122 input_image=input_image, 

123 model=model, 

124 size_i=size_i, 

125 size_j=size_j, 

126 size_k=size_k, 

127 overlap_i=overlap_i, 

128 overlap_j=overlap_j, 

129 overlap_k=overlap_k, 

130 spatial_dims=spatial_dims, 

131 device=device, 

132 single_crop=single_crop, 

133 ) 

134 write_image(prediction, output_path) 

135 

136 def generate_prediction( 

137 self, 

138 input_image: "torch.Tensor", 

139 model: "torch.nn.Module", 

140 size_i: int, 

141 size_j: int, 

142 size_k: int, 

143 overlap_i: int, 

144 overlap_j: int, 

145 overlap_k: int, 

146 spatial_dims: int, 

147 device: str, 

148 single_crop: bool, 

149 ) -> "torch.Tensor": 

150 import torch 

151 import math 

152 

153 from .utils import generate_overlapping_intervals, distance_to_boundary 

154 

155 episilon = 0.01 # small number so that we do not have a zero weight 

156 weight = episilon + distance_to_boundary( 

157 size_i=size_i, size_j=size_j, size_k=size_k if spatial_dims == 3 else 1 

158 ) 

159 if spatial_dims == 2: 

160 weight = weight[:, :, 0] 

161 

162 dtype = torch.float32 

163 

164 prediction = torch.zeros_like(input_image, dtype=dtype) 

165 summed_weights = torch.zeros_like(input_image, dtype=dtype) 

166 input_image = input_image.to(dtype=dtype).unsqueeze(0) # add batch dimension 

167 

168 intervals_i = generate_overlapping_intervals( 

169 prediction.shape[-spatial_dims], size_i, overlap_i 

170 ) 

171 intervals_j = generate_overlapping_intervals( 

172 prediction.shape[1 - spatial_dims], size_j, overlap_j 

173 ) 

174 if spatial_dims == 3: 

175 intervals_k = generate_overlapping_intervals( 

176 prediction.shape[2 - spatial_dims], size_k, overlap_k 

177 ) 

178 else: 

179 intervals_k = [(0, None)] 

180 

181 if single_crop: 

182 # choose the center crop only 

183 intervals_i = [intervals_i[len(intervals_i) // 2]] 

184 intervals_j = [intervals_j[len(intervals_j) // 2]] 

185 

186 if spatial_dims == 3: 

187 intervals_k = [intervals_k[len(intervals_k) // 2]] 

188 

189 total_tiles = len(intervals_i) * len(intervals_j) * len(intervals_k) 

190 progress_context = Progress() if total_tiles else nullcontext() 

191 

192 with torch.no_grad(): 

193 with progress_context as progress_bar: 

194 task_id = None 

195 if total_tiles: 

196 task_id = progress_bar.add_task( 

197 "Predicting tiles", total=total_tiles 

198 ) 

199 

200 for start_i, end_i in intervals_i: 

201 for start_j, end_j in intervals_j: 

202 for start_k, end_k in intervals_k: 

203 spatial_ranges = [ 

204 slice(start_i, end_i), 

205 slice(start_j, end_j), 

206 ] 

207 if spatial_dims == 3: 

208 spatial_ranges.append(slice(start_k, end_k)) 

209 

210 spatial_ranges = tuple(spatial_ranges) 

211 

212 cropped = input_image[ 

213 (slice(None), slice(None), *spatial_ranges) 

214 ] 

215 if torch.isnan(cropped).all(): 

216 continue 

217 

218 cropped = cropped.to(device=device) 

219 

220 result = model(cropped) 

221 prediction[spatial_ranges] += ( 

222 result[0, 0].to("cpu") * weight 

223 ) 

224 summed_weights[spatial_ranges] += weight 

225 

226 if task_id is not None: 

227 progress_bar.advance(task_id) 

228 

229 non_zero_voxels = summed_weights > 0 

230 prediction[non_zero_voxels] /= summed_weights[non_zero_voxels] 

231 prediction[~non_zero_voxels] = math.nan 

232 

233 return prediction.squeeze(0) # remove batch dimension 

234 

235 @tool 

236 def porosity( 

237 self, 

238 input: Path = cluey.Option(None, help="Path to the input image"), 

239 ): 

240 """Calculate and print the porosity of an input image.""" 

241 from .data import read_image_as_tensor 

242 from .metrics import calc_porosity 

243 

244 input_image = read_image_as_tensor(input) 

245 porosity = calc_porosity(input_image) 

246 print(f"Porosity: {porosity}") 

247 return porosity 

248 

249 @tool 

250 def porosity_distribution( 

251 self, 

252 input: Path = cluey.Option(None, help="Path to the input image"), 

253 output: Path = cluey.Option( 

254 None, help="Path to the output CSV containing seeds and porosities" 

255 ), 

256 size: int = cluey.Option( 

257 100, help="Default prediction tile size for all spatial dimensions" 

258 ), 

259 size_i: int = cluey.Option( 

260 0, help="Prediction tile size along the i dimension (0 uses size)" 

261 ), 

262 size_j: int = cluey.Option( 

263 0, help="Prediction tile size along the j dimension (0 uses size)" 

264 ), 

265 size_k: int = cluey.Option( 

266 0, help="Prediction tile size along the k dimension (0 uses size)" 

267 ), 

268 overlap: int = cluey.Option( 

269 10, help="Default overlap between prediction tiles in pixels or voxels" 

270 ), 

271 overlap_i: int = cluey.Option( 

272 0, help="Tile overlap along the i dimension (0 uses overlap)" 

273 ), 

274 overlap_j: int = cluey.Option( 

275 0, help="Tile overlap along the j dimension (0 uses overlap)" 

276 ), 

277 overlap_k: int = cluey.Option( 

278 0, help="Tile overlap along the k dimension (0 uses overlap)" 

279 ), 

280 checkpoint: Path = cluey.Option(None, help="Path to the model checkpoint"), 

281 num_sampling_steps: int = cluey.Option( 

282 250, help="Number of diffusion sampling steps per prediction" 

283 ), 

284 seed: int = cluey.Option( 

285 42, help="Starting random seed for consecutive diffusion samples" 

286 ), 

287 single_crop: bool = cluey.Option(False, help="Predict only the center tile"), 

288 count: int = cluey.Option( 

289 50, 

290 help="Number of consecutive seeds to sample, skipping seeds already saved", 

291 ), 

292 overwrite: bool = cluey.Option( 

293 False, help="Replace the existing output CSV instead of appending results" 

294 ), 

295 **kwargs, 

296 ): 

297 """Sample diffusion predictions and save their seeds and porosities to a CSV file.""" 

298 import torch 

299 from widit import load_model 

300 

301 from .data import read_image_as_tensor 

302 from .models import DiffusionPredictionModel 

303 from .metrics import calc_porosity 

304 

305 torch.set_grad_enabled(False) 

306 device = "cuda" if torch.cuda.is_available() else "cpu" 

307 

308 assert input is not None, "Must provide input path" 

309 input_image = read_image_as_tensor(input) 

310 spatial_dims = input_image.ndim - 1 

311 assert spatial_dims in ( 

312 2, 

313 3, 

314 ), f"Input image must have 3 or 4 dimensions (C, [D], H, W), got {input_image.shape}" 

315 

316 assert output is not None, "Must provide output path" 

317 output_path = Path(output) 

318 output_path.parent.mkdir(parents=True, exist_ok=True) 

319 

320 assert checkpoint is not None, "Must provide path to model checkpoint" 

321 

322 size_i = size_i or size 

323 size_j = size_j or size 

324 size_k = size_k or size 

325 

326 overlap_i = overlap_i or overlap 

327 overlap_j = overlap_j or overlap 

328 overlap_k = overlap_k or overlap 

329 

330 # TODO scale the image to (size_k, size_j, size_i) if it is not already that size 

331 

332 assert ( 

333 input_image.shape[-1] == size_i 

334 ), f"Input image size in i dimension ({input_image.shape[-1]}) does not match specified size_i ({size_i})" 

335 assert ( 

336 input_image.shape[-2] == size_j 

337 ), f"Input image size in j dimension ({input_image.shape[-2]}) does not match specified size_j ({size_j})" 

338 if spatial_dims == 3: 

339 assert ( 

340 input_image.shape[-3] == size_k 

341 ), f"Input image size in k dimension ({input_image.shape[-3]}) does not match specified size_k ({size_k})" 

342 

343 # Load checkpoint 

344 model = load_model(checkpoint) 

345 

346 diffusion = model.out_channels == 2 

347 assert ( 

348 diffusion 

349 ), "Model must be a diffusion model with 2 output channels for porosity distribution tool" 

350 if diffusion: 

351 model = DiffusionPredictionModel(model, num_sampling_steps) 

352 

353 model.to(device=device) 

354 model.eval() 

355 

356 if overwrite and output_path.exists(): 

357 output_path.unlink() 

358 

359 existing_seeds = set() 

360 write_header = True 

361 if output_path.exists(): 

362 write_header = False 

363 with open(output_path) as existing_file: 

364 for line in existing_file: 

365 line = line.strip() 

366 if not line or line == "seed,porosity": 

367 continue 

368 seed_value, _, _ = line.partition(",") 

369 try: 

370 existing_seeds.add(int(seed_value)) 

371 except ValueError: 

372 continue 

373 

374 with open(output_path, "a") as f: 

375 if write_header: 

376 f.write("seed,porosity\n") 

377 f.flush() 

378 

379 for index in range(count): 

380 current_seed = seed + index 

381 if current_seed in existing_seeds: 

382 print( 

383 f"Skipping seed {current_seed}; already present in {output_path}" 

384 ) 

385 continue 

386 

387 torch.manual_seed(current_seed) 

388 prediction = self.generate_prediction( 

389 input_image=input_image, 

390 model=model, 

391 size_i=size_i, 

392 size_j=size_j, 

393 size_k=size_k, 

394 overlap_i=overlap_i, 

395 overlap_j=overlap_j, 

396 overlap_k=overlap_k, 

397 spatial_dims=spatial_dims, 

398 device=device, 

399 single_crop=single_crop, 

400 ) 

401 porosity = calc_porosity(prediction) 

402 print(f"Seed: {current_seed}, Porosity: {porosity}") 

403 f.write(f"{current_seed},{porosity}\n") 

404 f.flush()