(
batch: Batch,
cfg: CroppingCfg,
)
| 94 | |
| 95 | |
| 96 | def crop_and_resize_batch_for_model( |
| 97 | batch: Batch, |
| 98 | cfg: CroppingCfg, |
| 99 | ) -> tuple[Batch, tuple[int, int]]: |
| 100 | # Resize the batch to the desired model input size. |
| 101 | image_shape = get_image_shape(tuple(batch.videos.shape[-2:]), cfg) |
| 102 | batch = resize_batch(batch, image_shape) |
| 103 | |
| 104 | # Record the pre-cropping shape. |
| 105 | _, _, _, h, w = batch.videos.shape |
| 106 | |
| 107 | # Center-crop the batch so it's cleanly divisible by the patch size. |
| 108 | return patch_crop_batch(batch, cfg.patch_size), (h, w) |
| 109 | |
| 110 | |
| 111 | def crop_and_resize_batch_for_flow(batch: Batch, cfg: CroppingCfg) -> Batch: |
no test coverage detected