Function crop_and_resize_batch_for_flow

flowmap/misc/cropping.py:111–122

(batch: Batch, cfg: CroppingCfg)

def crop_and_resize_batch_for_flow(batch: Batch, cfg: CroppingCfg) -> Batch:
    # Figure out the image size that's used for flow.
    image_shape = get_image_shape(tuple(batch.videos.shape[-2:]), cfg)
    flow_shape = tuple(dim * cfg.flow_scale_multiplier for dim in image_shape)

    # Resize the batch to match the desired flow shape.
    batch = resize_batch(batch, flow_shape)

    # Center-crop the batch so it's cleanly divisible by the patch size times the flow
    # multiplier. This ensures that the aspect ratio matches the model input's aspect
    # ratio.
    return patch_crop_batch(batch, cfg.patch_size * cfg.flow_scale_multiplier)


def resize_to_cover(