54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
import math
|
|
import random
|
|
|
|
import torch
|
|
from torchvision import transforms
|
|
from torchvision.transforms import Lambda
|
|
|
|
|
|
|
|
class TemporalRandomCrop(object):
|
|
"""Temporally crop the given frame indices at a random location.
|
|
|
|
Args:
|
|
size (int): Desired length of frames will be seen in the model.
|
|
"""
|
|
|
|
def __init__(self, size):
|
|
self.size = size
|
|
|
|
def __call__(self, total_frames):
|
|
rand_end = max(0, total_frames - self.size - 1)
|
|
begin_index = random.randint(0, rand_end)
|
|
end_index = min(begin_index + self.size, total_frames)
|
|
return begin_index, end_index
|
|
|
|
|
|
class LongSideScale(torch.nn.Module):
|
|
"""
|
|
``nn.Module`` wrapper for ``pytorchvideo.transforms.functional.short_side_scale``.
|
|
"""
|
|
|
|
def __init__(
|
|
self, size: int, interpolation: str = "bilinear"
|
|
):
|
|
super().__init__()
|
|
self._size = size
|
|
self._interpolation = interpolation
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x (torch.Tensor): video tensor with shape (C, T, H, W).
|
|
"""
|
|
assert len(x.shape) == 4
|
|
assert x.dtype == torch.float32
|
|
c, t, h, w = x.shape
|
|
if w < h:
|
|
new_w = int(math.floor((float(w) / h) * self._size))
|
|
new_h = self._size
|
|
else:
|
|
new_w = self._size
|
|
new_h = int(math.floor((float(h) / w) * self._size))
|
|
return torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=self._interpolation, align_corners=False)
|