| 430 | ) |
| 431 | |
| 432 | def augment_batch(batch): |
| 433 | # move everything to GPU |
| 434 | frames = batch["frames"] # [B, T, 3, H, W] |
| 435 | depths = batch["depths"] # [B, T, 1, H, W] |
| 436 | masks = batch["supv_masks"] # [B, T, 1, H, W] |
| 437 | |
| 438 | # map frames to [-1, 1] |
| 439 | frames = frames * 2 - 1 |
| 440 | |
| 441 | B, T, C, H, W = frames.shape |
| 442 | _, _, D, _, _ = depths.shape |
| 443 | _, _, M, _, _ = masks.shape |
| 444 | |
| 445 | # flatten the (B, T, ...) → (B*T, ...) |
| 446 | frames_flat = frames.view(-1, C, H, W) |
| 447 | depths_flat = depths.view(-1, D, H, W) |
| 448 | masks_flat = masks.view(-1, M, H, W) |
| 449 | |
| 450 | # 1) sample one set of geom_params per sequence (using the first frame of each seq) |
| 451 | seed = frames[:, 0] # [B, 3, H, W] |
| 452 | _, _, geom_params = geom_pipe(seed, return_params=True) |
| 453 | # now expand those params to all T frames in each sequence |
| 454 | geom_params = {k: v.repeat_interleave(T, dim=0) for k, v in geom_params.items()} |
| 455 | |
| 456 | # 2) apply the SAME geometry warp to every frame/depth/mask |
| 457 | frames_aug_flat, _, _ = geom_pipe(frames_flat, params=geom_params) |
| 458 | depths_aug_flat, _, _ = geom_pipe(depths_flat, params=geom_params) |
| 459 | masks_aug_flat, _, _ = geom_pipe(masks_flat, params=geom_params) |
| 460 | |
| 461 | # 3) reshape back to (B, T, …) |
| 462 | batch["frames"] = frames_aug_flat.view(B, T, C, H, W) |
| 463 | batch["depths"] = depths_aug_flat.view(B, T, D, H, W) |
| 464 | batch["supv_masks"] = masks_aug_flat.view(B, T, M, H, W) |
| 465 | |
| 466 | # (optionally) do the photo‐jittering in the same way: |
| 467 | seed_photo = batch["frames"][:, 0] |
| 468 | _, _, photo_params = photo_pipe(seed_photo, return_params=True) |
| 469 | photo_params = {k: v.repeat_interleave(T, dim=0) for k,v in photo_params.items()} |
| 470 | flat = batch["frames"].view(-1, C, H, W) |
| 471 | aug, _, _ = photo_pipe(flat, params=photo_params) |
| 472 | batch["frames"] = aug.view(B, T, C, H, W) |
| 473 | |
| 474 | # unmap frames back to [0, 1] |
| 475 | batch["frames"] = ((batch["frames"] + 1) / 2).clamp(0, 1) |
| 476 | |
| 477 | return batch |