Coverage for supercat/data.py: 100.00%

1from pathlib import Path 

2import numpy as np 

3import torch 

4from torch.utils.data import Dataset 

5from skimage.transform import rescale, resize 

6import hdf5storage 

7from PIL import Image 

8 

9 

10def read_image(path: str | Path, size: tuple[int, int, int] | None = None): 

11 path = Path(path) 

12 suffix = path.suffix.lower() 

13 if suffix in [".png", ".jpg", ".jpeg"]: 

14 img = Image.open(path).convert("L") 

15 if size is not None: 

16 size = size[:2] 

17 img = img.resize(size[::-1], Image.BICUBIC) 

18 elif suffix in [".mat"]: 

19 img = read_mat(path) 

20 assert img.ndim == 3, f"Expected 3D array from {path}, got shape {img.shape}" 

21 if size is not None: 

22 img = resize( 

23 img, 

24 output_shape=size, 

25 order=3, # cubic 

26 mode="reflect", 

27 anti_aliasing=False, # no AA on upscaling 

28 preserve_range=True, 

29 ) 

30 else: 

31 raise ValueError(f"Unsupported image format: {suffix}") 

32 

33 array = np.array(img, dtype=np.float32) 

34 return np.expand_dims(array, axis=0) 

35 

36 

37def transform_scale(data): 

38 return 2.0 * data / 255.0 - 1 

39 

40 

41def read_image_as_tensor(path, **kwargs) -> torch.Tensor: 

42 image = read_image(path, **kwargs) 

43 image = transform_scale(image) 

44 

45 image = torch.from_numpy(image.copy()) # ensure contiguous 

46 

47 return image 

48 

49 

50TRANSFORMATIONS_3D = [ 

51 lambda tensor: tensor, # Identity 

52 lambda tensor: tensor.flip(-2).transpose(-1, -2), # 90° rotation 

53 lambda tensor: tensor.flip(-2).flip(-1), # 180° rotation 

54 lambda tensor: tensor.flip(-2) 

55 .transpose(-1, -2) 

56 .flip(-1), # TR-BL diagonal reflection 

57 lambda tensor: tensor.flip(-1), # Vertical reflection 

58 lambda tensor: tensor.flip(-2), # Horizontal reflection 

59 lambda tensor: tensor.flip(-1).transpose(-1, -2), # TL-BR diagonal reflection 

60 lambda tensor: tensor.flip(-1) 

61 .flip(-2) 

62 .transpose(-1, -2), # TR-BL diagonal reflection with both flips 

63 lambda tensor: tensor.flip(-3), 

64 lambda tensor: tensor.flip(-1).flip(-3), 

65 lambda tensor: tensor.flip(-2).flip(-3), 

66 lambda tensor: tensor.flip(-1).flip(-2).flip(-3), 

67 lambda tensor: tensor.transpose(-1, -2).flip(-3), 

68 lambda tensor: tensor.transpose(-1, -2).flip(-1).flip(-3), 

69 lambda tensor: tensor.transpose(-1, -2).flip(-2).flip(-3), 

70 lambda tensor: tensor.transpose(-1, -2).flip(-1).flip(-2).flip(-3), 

71 lambda tensor: tensor.transpose(-1, -3), 

72 lambda tensor: tensor.transpose(-1, -3).flip(-1), 

73 lambda tensor: tensor.transpose(-1, -3).flip(-2), 

74 lambda tensor: tensor.transpose(-1, -3).flip(-3), 

75 lambda tensor: tensor.transpose(-1, -3).flip(-1).flip(-2), 

76 lambda tensor: tensor.transpose(-1, -3).flip(-1).flip(-3), 

77 lambda tensor: tensor.transpose(-1, -3).flip(-2).flip(-3), 

78 lambda tensor: tensor.transpose(-1, -3).flip(-1).flip(-2).flip(-3), 

79 lambda tensor: tensor.transpose(-2, -3), 

80 lambda tensor: tensor.transpose(-2, -3).flip(-1), 

81 lambda tensor: tensor.transpose(-2, -3).flip(-2), 

82 lambda tensor: tensor.transpose(-2, -3).flip(-3), 

83 lambda tensor: tensor.transpose(-2, -3).flip(-1).flip(-2), 

84 lambda tensor: tensor.transpose(-2, -3).flip(-1).flip(-3), 

85 lambda tensor: tensor.transpose(-2, -3).flip(-2).flip(-3), 

86 lambda tensor: tensor.transpose(-2, -3).flip(-1).flip(-2).flip(-3), 

87 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2), 

88 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-1), 

89 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-2), 

90 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-3), 

91 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-1).flip(-2), 

92 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-1).flip(-3), 

93 lambda tensor: tensor.transpose(-2, -3).transpose(-1, -2).flip(-2).flip(-3), 

94 lambda tensor: tensor.transpose(-2, -3) 

95 .transpose(-1, -2) 

96 .flip(-1) 

97 .flip(-2) 

98 .flip(-3), 

99 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3), 

100 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-1), 

101 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-2), 

102 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-3), 

103 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-1).flip(-2), 

104 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-1).flip(-3), 

105 lambda tensor: tensor.transpose(-1, -2).transpose(-2, -3).flip(-2).flip(-3), 

106 lambda tensor: tensor.transpose(-1, -2) 

107 .transpose(-2, -3) 

108 .flip(-1) 

109 .flip(-2) 

110 .flip(-3), 

111] 

112TRANSFORMATIONS_2D = TRANSFORMATIONS_3D[:8] 

113 

114# def flip_and_rotate( 

115# batch:tuple[torch.Tensor,torch.Tensor,torch.Tensor], 

116# sym_id:int|None=None 

117# ) -> tuple[torch.Tensor,torch.Tensor,torch.Tensor]: 

118# """ 

119# Applies a random or specified flip and/or rotation transformation to a batch of 2D or 3D tensors. 

120# This function transforms the input, target, and residual tensors by flipping, rotating, and 

121# transposing them in a consistent manner. 

122 

123# Args: 

124# batch (tuple[torch.Tensor, torch.Tensor, torch.Tensor]): A tuple containing the input, target, 

125# and residual tensors to be transformed. 

126# sym_id (int or None, optional): The transformation identifier. If None, a random transformation 

127# is chosen. For 2D tensors, sym_id ranges from 0 to 7; for 3D tensors, it ranges from 0 to 47. 

128 

129# Returns: 

130# tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing the transformed input, 

131# target, and residual tensors after applying the flip and/or rotation. 

132 

133# Raises: 

134# AssertionError: If the sym_id is out of the valid range for the tensor dimensions. 

135# """ 

136# input, target, residual = batch 

137 

138# # Determine if 2D (4D) or 3D (5D) 

139# is_2d = (input.ndim == 4) 

140# number_of_transforms = 8 if is_2d else len(TRANSFORMATIONS) 

141 

142# if sym_id is None: 

143# # Randomly choose one of the transformations 

144# sym_id = random.randint(0, number_of_transforms - 1) 

145 

146# # Shortcut for no transformation 

147# if sym_id == 0: 

148# return batch 

149 

150# assert 0 <= sym_id < number_of_transforms 

151# transformation = TRANSFORMATIONS[sym_id] 

152 

153# input = transformation(input) 

154# target = transformation(target) 

155# residual = transformation(residual) 

156 

157# return input, target, residual 

158 

159 

160def downscale_tricubic_rescale(vol: np.ndarray, factor: int = 4) -> np.ndarray: 

161 """ 

162 Downscale a 3D volume by `factor` using tricubic interpolation + anti-aliasing. 

163 """ 

164 low = rescale( 

165 vol, 

166 scale=1.0 / factor, 

167 order=3, # cubic 

168 mode="reflect", 

169 anti_aliasing=True, 

170 preserve_range=True, 

171 channel_axis=None, # ensures (D,H,W) or (C,D,H,W) handled correctly if needed 

172 ) 

173 return low.astype(vol.dtype, copy=False) 

174 

175 

176def upscale_tricubic_rescale(vol: np.ndarray, factor: int = 4) -> np.ndarray: 

177 """ 

178 Upscale a 3D volume by `factor` using tricubic interpolation. 

179 """ 

180 up = rescale( 

181 vol, 

182 scale=factor, 

183 order=3, # cubic 

184 mode="reflect", 

185 anti_aliasing=False, # no AA on upscaling 

186 preserve_range=True, 

187 channel_axis=None, 

188 ) 

189 return up.astype(vol.dtype, copy=False) 

190 

191 

192def read_mat(path: Path): 

193 DEEPROCK_HDF5_KEY = "temp" 

194 try: 

195 data_dict = hdf5storage.loadmat(str(path)) 

196 except Exception as err: 

197 raise IOError(f"Error reading 3D file '{path}':\n{err}") 

198 if DEEPROCK_HDF5_KEY not in data_dict: 

199 keys_found = ",".join(data_dict.keys()) 

200 raise Exception( 

201 f"expected key {DEEPROCK_HDF5_KEY} not found in '{path}'.\nCheck the following keys: {keys_found}" 

202 ) 

203 

204 return data_dict[DEEPROCK_HDF5_KEY] 

205 

206 

207class Deeprock3D(Dataset): 

208 def __init__( 

209 self, 

210 deeprock: Path, 

211 scale: int = 4, 

212 channel_first: bool = True, 

213 partition: str = "train", 

214 augment: bool = False, 

215 ): 

216 self.deeprock = Path(deeprock) 

217 self.scale = int(scale) 

218 self.channel_first = channel_first 

219 categories = ["sandstone", "carbonate", "coal", "sand"] 

220 self.augment = augment 

221 

222 hr_items: list[Path] = [] 

223 for cat in categories: 

224 hr_dir = self.deeprock / f"{cat}3D" / f"{cat}3D_{partition}_HR" 

225 hr_items.extend(sorted(hr_dir.glob("*.mat"))) 

226 self.hr_items = hr_items 

227 

228 if len(self.hr_items) == 0: 

229 raise FileNotFoundError( 

230 f"No .mat files found under {self.deeprock}/*3D/*3D_{partition}_HR" 

231 ) 

232 

233 def _hr_to_lr_path(self, hr_path: Path) -> Path: 

234 """ 

235 Map .../<cat>3D/<cat>3D_train_HR/<file>.mat --> 

236 .../<cat>3D/<cat>3D_train_TRI_unknown_X{scale}/<file>.mat 

237 """ 

238 rel = hr_path.relative_to(self.deeprock) 

239 parts = list(rel.parts) 

240 # parts[-2] expected: "<cat>3D_train_HR" 

241 parent_name = parts[-2] 

242 if not parent_name.endswith("_HR"): 

243 raise ValueError(f"Unexpected parent '{parent_name}' for {hr_path}") 

244 parts[-2] = parent_name.replace("_HR", f"_LR_default_X{self.scale}").replace( 

245 ".mat", f"x{self.scale}.mat" 

246 ) 

247 parts[-1] = parts[-1].replace(".mat", f"x{self.scale}.mat") 

248 return self.deeprock.joinpath(*parts) 

249 

250 def __len__(self) -> int: 

251 return len(self.hr_items) 

252 

253 def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: 

254 hr_path = self.hr_items[idx] 

255 

256 lr_path = self._hr_to_lr_path(hr_path) 

257 

258 # if not lr_path.exists(): 

259 # raise FileNotFoundError(f"LR file missing for {hr_path.name}: {lr_path}") 

260 

261 hr = transform_scale(read_mat(hr_path)) 

262 lr_orig = transform_scale(read_mat(lr_path)) 

263 # lr_orig = downscale_tricubic_rescale(lr_orig, factor=self.scale) 

264 lr = upscale_tricubic_rescale(lr_orig, factor=self.scale) 

265 

266 # lr = transform_scale(read_mat(lr_path)) 

267 

268 # Optional: enforce channel-first tensors 

269 # If (D,H,W), add channel dim; if already (C,D,H,W), leave it. 

270 if hr.ndim == 3: 

271 hr = hr[None, ...] # (1,D,H,W) 

272 if lr.ndim == 3: 

273 lr = lr[None, ...] 

274 if not self.channel_first and hr.ndim == 4: 

275 # Convert to (D,H,W,C) 

276 hr = np.moveaxis(hr, 0, -1) 

277 lr = np.moveaxis(lr, 0, -1) 

278 

279 hr_t = torch.from_numpy(hr.copy()) # ensure contiguous 

280 lr_t = torch.from_numpy(lr.copy()) 

281 

282 if self.augment: 

283 transformation = TRANSFORMATIONS_3D[ 

284 np.random.randint(0, len(TRANSFORMATIONS_3D)) 

285 ] 

286 hr_t = transformation(hr_t) 

287 lr_t = transformation(lr_t) 

288 

289 assert ( 

290 lr_t.shape == hr_t.shape 

291 ), f"LR shape {lr_t.shape} != HR shape {hr_t.shape} for {hr_path.name} - {lr_path.name}, original LR shape {lr_orig.shape}. scale={self.scale}" 

292 assert ( 

293 lr_t.max() < 1.01 

294 ), f"Low resolution {lr_path} gives range {lr_t.min()}-{lr_t.max()}" 

295 assert ( 

296 lr_t.min() > -1.01 

297 ), f"Low resolution {lr_path} gives range {lr_t.min()}-{lr_t.max()}" 

298 assert ( 

299 hr_t.max() < 1.01 

300 ), f"High resolution {hr_path} gives range {hr_t.min()}-{hr_t.max()}" 

301 assert ( 

302 hr_t.min() > -1.01 

303 ), f"High resolution {hr_path} gives range {hr_t.min()}-{hr_t.max()}" 

304 

305 return lr_t, hr_t 

306 

307 

308class Deeprock2D(Dataset): 

309 def __init__( 

310 self, 

311 deeprock: Path, 

312 scale: int = 4, 

313 channel_first: bool = True, 

314 partition: str = "train", 

315 augment: bool = False, 

316 ): 

317 self.deeprock = Path(deeprock) 

318 self.scale = int(scale) 

319 self.channel_first = channel_first 

320 categories = ["sandstone", "carbonate", "coal", "sand"] 

321 self.augment = augment 

322 

323 hr_items: list[Path] = [] 

324 for cat in categories: 

325 hr_dir = self.deeprock / f"{cat}2D" / f"{cat}2D_{partition}_HR" 

326 hr_items.extend(sorted(hr_dir.glob("*.png"))) 

327 self.hr_items = hr_items 

328 

329 if len(self.hr_items) == 0: 

330 raise FileNotFoundError( 

331 f"No .png files found under {self.deeprock}/*2D/*2D_{partition}_HR" 

332 ) 

333 

334 def _hr_to_lr_path(self, hr_path: Path) -> Path: 

335 """ 

336 Map .../<cat>2D/<cat>2D_train_HR/<file>.png --> 

337 .../<cat>2D/<cat>2D_train_BI_unknown_X{scale}/<file>.png 

338 """ 

339 rel = hr_path.relative_to(self.deeprock) 

340 parts = list(rel.parts) 

341 # parts[-2] expected: "<cat>2D_train_HR" 

342 parent_name = parts[-2] 

343 if not parent_name.endswith("_HR"): 

344 raise ValueError(f"Unexpected parent '{parent_name}' for {hr_path}") 

345 parts[-2] = parent_name.replace("_HR", f"_BI_unknown_X{self.scale}") 

346 return self.deeprock.joinpath(*parts) 

347 

348 def __len__(self) -> int: 

349 return len(self.hr_items) 

350 

351 def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: 

352 hr_path = self.hr_items[idx] 

353 lr_path = self._hr_to_lr_path(hr_path) 

354 

355 if not lr_path.exists(): 

356 raise FileNotFoundError(f"LR file missing for {hr_path.name}: {lr_path}") 

357 

358 hr_t = read_image_as_tensor(hr_path) 

359 lr_t = read_image_as_tensor(lr_path) 

360 

361 assert hr_t.shape[0] == 1 

362 

363 if self.augment: 

364 transformation = TRANSFORMATIONS_2D[ 

365 np.random.randint(0, len(TRANSFORMATIONS_2D)) 

366 ] 

367 hr_t = transformation(hr_t) 

368 lr_t = transformation(lr_t) 

369 

370 return lr_t, hr_t 

371 

372 

373def build_datasets3D(deeprock: Path, scale: int = 4, train_augment: bool = True): 

374 training_dataset = Deeprock3D( 

375 deeprock=deeprock, scale=scale, partition="train", augment=train_augment 

376 ) 

377 validation_dataset = Deeprock3D( 

378 deeprock=deeprock, scale=scale, partition="valid", augment=False 

379 ) 

380 return training_dataset, validation_dataset 

381 

382 

383def build_datasets2D(deeprock: Path, scale: int = 4, train_augment: bool = True): 

384 training_dataset = Deeprock2D( 

385 deeprock=deeprock, scale=scale, partition="train", augment=train_augment 

386 ) 

387 validation_dataset = Deeprock2D( 

388 deeprock=deeprock, scale=scale, partition="valid", augment=False 

389 ) 

390 return training_dataset, validation_dataset