| from typing import Callable, Any | |
| from torch import Tensor | |
| from einops import rearrange | |
| def videos_as_images( | |
| func: Callable[..., Tensor], num_video_args: int = 1 | |
| ) -> Callable[..., Tensor]: | |
| """ | |
| Wrapper that enables a function that operates on a batch of images to operate on a batch of videos. | |
| Can also be used as a decorator. | |
| """ | |
| def wrapper(*args: Any, **kwargs: Any) -> Tensor: | |
| # check if the first argument is a tensor or not | |
| # if not, assume it is "self" | |
| new_args = list(args) | |
| videos_idx = 1 if not isinstance(new_args[0], Tensor) else 0 | |
| b = new_args[videos_idx].shape[0] | |
| for idx in range(videos_idx, videos_idx + num_video_args): | |
| new_args[idx] = rearrange(new_args[idx], "b t c h w -> (b t) c h w") | |
| return rearrange( | |
| func(*new_args, **kwargs), | |
| "(b t) ... -> b t ...", | |
| b=b, | |
| ) | |
| return wrapper | |