Coverage for supercat/pretrain.py: 100.00%

1import functools

2import math

3from pathlib import Path

4import random

5import re

6import shutil

7import subprocess

9import numpy as np

10import torch

11import torch.nn.functional as F

12import cluey

13from cluey import method

14from torch.utils.data import Dataset

15from widitapp import WiDiTApp

17from .data import read_image_as_tensor, TRANSFORMATIONS_2D, TRANSFORMATIONS_3D

20def _normalize_video_frame(frame: np.ndarray) -> np.ndarray:

21 """Convert a decoded frame to a 2D grayscale array."""

22 if frame.ndim == 2:

23 return frame

24 if frame.ndim == 3 and frame.shape[-1] == 1:

25 return frame[..., 0]

26 if frame.ndim == 3:

27 # Match PIL's default grayscale conversion closely enough for training.

28 return np.round(frame[..., :3].mean(axis=-1)).astype(frame.dtype)

29 raise ValueError(f"Unsupported frame shape: {frame.shape}")

32def _ffmpeg_binary() -> str | None:

33 return shutil.which("ffmpeg")

36def _probe_ffmpeg_video(path: Path) -> tuple[int, int, int]:

37 ffmpeg = _ffmpeg_binary()

38 if ffmpeg is None:

39 raise RuntimeError("ffmpeg binary not found")

41 proc = subprocess.run(

42 [ffmpeg, "-hide_banner", "-i", str(path), "-map", "0:v:0", "-f", "null", "-"],

43 capture_output=True,

44 text=True,

45 check=False,

46 )

47 output = proc.stdout + proc.stderr

49 size_match = re.search(r"Video:.*?,\s+(\d+)x(\d+)\s+\[", output)

50 if size_match is None:

51 raise RuntimeError(f"Could not determine video size for {path}")

52 frame_width = int(size_match.group(1))

53 frame_height = int(size_match.group(2))

55 frame_matches = re.findall(r"frame=\s*(\d+)", output)

56 if not frame_matches:

57 raise RuntimeError(f"Could not determine frame count for {path}")

58 frame_count = int(frame_matches[-1])

60 return frame_count, frame_height, frame_width

63@functools.lru_cache(maxsize=8)

64def _read_video_via_ffmpeg(path: str) -> np.ndarray:

65 video_path = Path(path)

66 frame_count, frame_height, frame_width = _probe_ffmpeg_video(video_path)

67 ffmpeg = _ffmpeg_binary()

68 if ffmpeg is None:

69 raise RuntimeError("ffmpeg binary not found")

71 proc = subprocess.run(

72 [

73 ffmpeg,

74 "-hide_banner",

75 "-loglevel",

76 "error",

77 "-i",

78 str(video_path),

79 "-map",

80 "0:v:0",

81 "-vf",

82 "format=gray",

83 "-f",

84 "rawvideo",

85 "-pix_fmt",

86 "gray",

87 "-",

88 ],

89 capture_output=True,

90 check=True,

91 )

93 video = np.frombuffer(proc.stdout, dtype=np.uint8)

94 frame_size = frame_height * frame_width

95 actual_frame_count, remainder = divmod(video.size, frame_size)

96 if remainder != 0:

97 raise RuntimeError(f"Unexpected ffmpeg output size for {video_path}")

99 if actual_frame_count != frame_count:

100 frame_count = actual_frame_count

101

102 return video.reshape(frame_count, frame_height, frame_width)

103

104

105def _video_backend(path: Path):

106 try:

107 import imageio.v3 as iio

108 except ImportError:

109 iio = None

110

111 if iio is not None:

112 try:

113 iio.improps(path)

114 return "imageio", iio

115 except OSError:

116 pass

117

118 try:

119 from skvideo.io import ffprobe, vreader

120 except ImportError:

121 ffprobe = None

122 vreader = None

123

124 if ffprobe is not None and vreader is not None:

125 try:

126 next(vreader(str(path), as_grey=True))

127 return "skvideo", (ffprobe, vreader)

128 except (AssertionError, OSError, StopIteration):

129 pass

130

131 if _ffmpeg_binary() is not None:

132 return "ffmpeg_cli", None

133

134 raise RuntimeError(

135 "No supported video backend is available. Install `imageio[ffmpeg]` "

136 "or make an `ffmpeg` binary available on PATH."

137 )

138

139

140def iter_video_frames(path: Path):

141 backend, reader_impl = _video_backend(path)

142 if backend == "imageio":

143 for frame in reader_impl.imiter(path):

144 yield _normalize_video_frame(np.asarray(frame))

145 return

146

147 if backend == "ffmpeg_cli":

148 for frame in _read_video_via_ffmpeg(str(path)):

149 yield frame

150 return

151

152 _, vreader = reader_impl

153 for frame in vreader(str(path), as_grey=True):

154 yield np.asarray(frame[0, :, :, 0])

155

156

157def video_shape(path: Path) -> tuple[int, int, int]:

158 backend, reader_impl = _video_backend(path)

159

160 if backend == "imageio":

161 props = reader_impl.improps(path)

162 shape = props.shape

163 if len(shape) < 3:

164 raise ValueError(

165 f"Unsupported video shape from imageio for {path}: {shape}"

166 )

167 frame_count, frame_height, frame_width = shape[:3]

168 if not math.isfinite(float(frame_count)):

169 frame_count = sum(1 for _ in reader_impl.imiter(path))

170 return int(frame_count), int(frame_height), int(frame_width)

171

172 if backend == "ffmpeg_cli":

173 video = _read_video_via_ffmpeg(str(path))

174 frame_count, frame_height, frame_width = video.shape

175 return int(frame_count), int(frame_height), int(frame_width)

176

177 ffprobe, _ = reader_impl

178 metadata = ffprobe(str(path))

179 video_metadata = metadata.get("video", {})

180

181 frame_height = int(video_metadata["@height"])

182 frame_width = int(video_metadata["@width"])

183

184 try:

185 frame_count = int(video_metadata["@nb_frames"])

186 except (KeyError, ValueError):

187 frame_count = sum(1 for _ in iter_video_frames(path))

188

189 return frame_count, frame_height, frame_width

190

191

192def _pad_to_size_with_reflect(

193 image: torch.Tensor,

194 target_size: int,

195 spatial_dims: int = 2,

196) -> torch.Tensor:

197 """

198 Pad the trailing spatial dimensions to ``target_size`` using safe reflect steps.

199

200 For each spatial dimension, pad only on the "right" side until that dimension

201 reaches ``target_size``. Uses reflect padding where possible, and falls back to

202 replicate padding when the size along a dimension is 1.

203

204 Examples

205 --------

206 2D expected input shape:

207 (..., H, W)

208

209 3D expected input shape:

210 (..., D, H, W)

211 """

212 if spatial_dims not in (2, 3):

213 raise ValueError(f"spatial_dims must be 2 or 3, got {spatial_dims}")

214

215 if image.ndim < spatial_dims:

216 raise ValueError(

217 f"image has {image.ndim} dims, but spatial_dims={spatial_dims}"

218 )

219

220 # Iterate from last spatial dim backwards:

221 # 2D: -1, -2

222 # 3D: -1, -2, -3

223 for dim in range(-1, -spatial_dims - 1, -1):

224 while image.shape[dim] < target_size:

225 remaining = target_size - image.shape[dim]

226 current = image.shape[dim]

227

228 if current <= 1:

229 # Reflection needs at least 2 pixels along the padded dimension.

230 mode = "replicate"

231 pad_amount = remaining

232 else:

233 mode = "reflect"

234 pad_amount = min(remaining, current - 1)

235

236 # F.pad expects padding specified from the last dimension moving left:

237 # 2D: (W_left, W_right, H_left, H_right)

238 # 3D: (W_left, W_right, H_left, H_right, D_left, D_right)

239 pad = [0] * (2 * spatial_dims)

240

241 # For dim=-1 -> index 1

242 # For dim=-2 -> index 3

243 # For dim=-3 -> index 5

244 pad_index = 2 * (-dim) - 1

245 pad[pad_index] = pad_amount

246

247 image = F.pad(image, tuple(pad), mode=mode)

248

249 return image

250

251

252def find_files(

253 base_path: Path, valid_extensions, max_items: int | None = None

254) -> list[Path]:

255 """Find all files under ``base_path`` matching any of the given extensions."""

256 if not base_path.is_dir():

257 raise ValueError(f"{base_path} is not a valid directory")

258

259 files = []

260 for ext in valid_extensions:

261 files.extend(base_path.rglob(f"*{ext.lower()}"))

262 files.extend(base_path.rglob(f"*{ext.upper()}"))

263

264 if max_items is not None and len(files) >= max_items:

265 return files[:max_items]

266

267 return files

268

269

270def find_images(base_path: Path) -> list[Path]:

271 return find_files(base_path, [".jpg", ".jpeg", ".png", ".bmp", ".tiff"])

272

273

274def find_movies(base_path: Path, max_items: int | None = None) -> list[Path]:

275 return find_files(base_path, [".mp4", ".avi", ".mov", ".mkv"], max_items=max_items)

276

277

278class PretrainImagesDataset(Dataset):

279 def __init__(

280 self,

281 path: Path,

282 scale: int = 4,

283 channel_first: bool = True,

284 augment: bool = False,

285 size: int = 0,

286 min_size: int = 0,

287 max_size: int = 0,

288 ):

289 assert path is not None, "Path must be provided"

290 self.path = Path(path)

291 self.scale = int(scale)

292 self.channel_first = channel_first

293 self.augment = augment

294 self.size = size

295 if size:

296 min_size = size

297 max_size = size

298 self.min_size = min_size

299 self.max_size = max_size

300

301 self.items = find_images(path)

302

303 def __len__(self) -> int:

304 return len(self.items)

305

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

307 path = self.items[idx]

308 hr_t = read_image_as_tensor(path)

309

310 if self.max_size:

311 if hr_t.shape[-1] > self.max_size:

312 start = (

313 np.random.randint(0, hr_t.shape[-1] - self.max_size)

314 if self.augment

315 else hr_t.shape[-1] // 2 - self.max_size // 2

316 )

317 hr_t = hr_t[..., start : start + self.max_size]

318 if hr_t.shape[-2] > self.max_size:

319 start = (

320 np.random.randint(0, hr_t.shape[-2] - self.max_size)

321 if self.augment

322 else hr_t.shape[-2] // 2 - self.max_size // 2

323 )

324 hr_t = hr_t[..., start : start + self.max_size, :]

325

326 if self.min_size:

327 hr_t = _pad_to_size_with_reflect(hr_t, self.min_size)

328

329 # Ensure dimensions are even

330 if hr_t.shape[-1] % 2 != 0:

331 hr_t = hr_t[..., :-1]

332 if hr_t.shape[-2] % 2 != 0:

333 hr_t = hr_t[..., :-1, :]

334

335 lr_t = torch.nn.functional.interpolate(

336 hr_t.unsqueeze(0),

337 scale_factor=1 / self.scale,

338 mode="bicubic",

339 align_corners=False,

340 ).squeeze(0)

341

342 lr_t = torch.nn.functional.interpolate(

343 lr_t.unsqueeze(0),

344 size=hr_t.shape[1:],

345 mode="bicubic",

346 align_corners=False,

347 ).squeeze(0)

348

349 if self.augment:

350 transformation = TRANSFORMATIONS_2D[

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

352 ]

353 hr_t = transformation(hr_t)

354 lr_t = transformation(lr_t)

355

356 return lr_t, hr_t

357

358

359class PretrainMovieDataset(Dataset):

360 def __init__(

361 self,

362 path: Path,

363 scale: int = 4,

364 channel_first: bool = True,

365 augment: bool = False,

366 size: int = 0,

367 min_size: int = 0,

368 max_size: int = 0,

369 max_items: int | None = None,

370 ):

371 assert path is not None, "Path must be provided"

372 self.path = Path(path)

373 self.scale = int(scale)

374 self.channel_first = channel_first

375 self.augment = augment

376 self.size = size

377 if size:

378 min_size = size

379 max_size = size

380 self.min_size = min_size

381 self.max_size = max_size

382

383 self.items = find_movies(path, max_items=max_items)

384

385 def __len__(self) -> int:

386 return len(self.items)

387

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

389 path = self.items[idx]

390

391 try:

392 frame_count, frame_height, frame_width = video_shape(path)

393

394 depth = self.max_size or frame_count

395 height = self.max_size or frame_height

396 width = self.max_size or frame_width

397

398 frame_start = (

399 random.randint(0, max(frame_count - depth, 0))

400 if self.augment

401 else max(frame_count // 2 - depth // 2, 0)

402 )

403 frame_end = min(frame_start + depth, frame_count)

404

405 y_start = (

406 random.randint(0, max(frame_height - height, 0))

407 if self.augment

408 else max(frame_height // 2 - height // 2, 0)

409 )

410 y_end = min(y_start + height, frame_height)

411 x_start = (

412 random.randint(0, max(frame_width - width, 0))

413 if self.augment

414 else max(frame_width // 2 - width // 2, 0)

415 )

416 x_end = min(x_start + width, frame_width)

417

418 image = np.zeros(

419 (frame_end - frame_start, y_end - y_start, x_end - x_start),

420 dtype=np.uint8,

421 )

422

423 for i, frame in enumerate(iter_video_frames(path)):

424 if i < frame_start:

425 continue

426 if i >= frame_end:

427 break

428 image[i - frame_start, :, :] = frame[y_start:y_end, x_start:x_end]

429

430 hr_t = torch.tensor(image / 255.0 * 2 - 1.0, dtype=torch.float32).unsqueeze(

431 0

432 )

433 except Exception as e:

434 print(f"Failed to load video {path}: {e}")

435 zeros = torch.zeros(

436 (1, self.min_size or 16, self.min_size or 16, self.min_size or 16),

437 dtype=torch.float32,

438 )

439 return zeros, zeros

440

441 # Ensure dimensions are even

442 if hr_t.shape[-1] % 2 != 0:

443 hr_t = hr_t[..., :-1]

444 if hr_t.shape[-2] % 2 != 0:

445 hr_t = hr_t[..., :-1, :]

446 if hr_t.shape[-3] % 2 != 0:

447 hr_t = hr_t[..., :-1, :, :]

448

449 # Pad to the target size if needed, since some videos may be smaller than the requested crop size.

450 if self.min_size:

451 hr_t = _pad_to_size_with_reflect(hr_t, self.min_size, spatial_dims=3)

452

453 lr_t = torch.nn.functional.interpolate(

454 hr_t.unsqueeze(0),

455 scale_factor=1 / self.scale,

456 mode="trilinear",

457 align_corners=False,

458 ).squeeze(0)

459

460 lr_t = torch.nn.functional.interpolate(

461 lr_t.unsqueeze(0),

462 size=hr_t.shape[1:],

463 mode="trilinear",

464 align_corners=False,

465 ).squeeze(0)

466

467 if self.augment:

468 transformation = TRANSFORMATIONS_3D[

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

470 ]

471 hr_t = transformation(hr_t)

472 lr_t = transformation(lr_t)

473

474 if not self.channel_first:

475 hr_t = hr_t.movedim(0, -1)

476 lr_t = lr_t.movedim(0, -1)

477

478 return lr_t, hr_t

479

480

481class SupercatPretrainImage(WiDiTApp):

482 @method

483 def datasets(

484 self,

485 training: Path = cluey.Option(

486 None, help="Path to the training image directory"

487 ),

488 validation: Path = cluey.Option(

489 None, help="Path to the validation image directory"

490 ),

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

492 augment: bool = cluey.Option(

493 True,

494 help="Apply random cropping and augmentation to the training image samples",

495 ),

496 min_size: int = cluey.Option(

497 16,

498 help="Minimum size for each spatial dimension after padding (0 disables padding)",

499 ),

500 max_size: int = cluey.Option(

501 224,

502 help="Maximum crop size for each spatial dimension (0 disables cropping)",

503 ),

504 **kwargs,

505 ) -> tuple[Dataset, Dataset]:

506 """Build training and validation datasets for 2D image pretraining."""

507 assert training is not None, "Training path must be provided"

508 assert validation is not None, "Validation path must be provided"

509 training_dataset = PretrainImagesDataset(

510 path=training,

511 scale=scale,

512 min_size=min_size,

513 max_size=max_size,

514 augment=augment,

515 )

516 validation_dataset = PretrainImagesDataset(

517 path=validation,

518 scale=scale,

519 min_size=min_size,

520 max_size=max_size,

521 augment=False,

522 )

523 return training_dataset, validation_dataset

524

525

526class SupercatPretrainMovie(WiDiTApp):

527 @method

528 def datasets(

529 self,

530 training: Path = cluey.Option(

531 None, help="Path to the training video directory"

532 ),

533 validation: Path = cluey.Option(

534 None, help="Path to the validation video directory"

535 ),

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

537 augment: bool = cluey.Option(

538 True,

539 help="Apply random cropping and augmentation to the training video samples",

540 ),

541 min_size: int = cluey.Option(

542 16,

543 help="Minimum frame count, height, and width after padding (0 disables padding)",

544 ),

545 max_size: int = cluey.Option(

546 100,

547 help="Maximum frame count, height, and width per crop (0 uses the full video)",

548 ),

549 max_training_items: int | None = cluey.Option(

550 None, help="Maximum number of training videos (None uses all videos)"

551 ),

552 max_validation_items: int | None = cluey.Option(

553 None, help="Maximum number of validation videos (None uses all videos)"

554 ),

555 **kwargs,

556 ) -> tuple[Dataset, Dataset]:

557 """Build training and validation datasets for 3D video pretraining."""

558 assert training is not None, "Training path must be provided"

559 assert validation is not None, "Validation path must be provided"

560

561 training_dataset = PretrainMovieDataset(

562 path=training,

563 scale=scale,

564 min_size=min_size,

565 max_size=max_size,

566 augment=augment,

567 max_items=max_training_items,

568 )

569 validation_dataset = PretrainMovieDataset(

570 path=validation,

571 scale=scale,

572 min_size=min_size,

573 max_size=max_size,

574 augment=False,

575 max_items=max_validation_items,

576 )

577 return training_dataset, validation_dataset