zamba.pytorch.transforms

ConvertHWCtoCHW

Bases: Module

Convert tensor from (0:H, 1:W, 2:C) to (2:C, 0:H, 1:W)

Source code in zamba/pytorch/transforms.py

class ConvertHWCtoCHW(torch.nn.Module):
    """Convert tensor from (0:H, 1:W, 2:C) to (2:C, 0:H, 1:W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(2, 0, 1)

ConvertTCHWtoCTHW

Bases: Module

Convert tensor from (T, C, H, W) to (C, T, H, W)

Source code in zamba/pytorch/transforms.py

class ConvertTCHWtoCTHW(torch.nn.Module):
    """Convert tensor from (T, C, H, W) to (C, T, H, W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(1, 0, 2, 3)

ConvertTHWCtoCTHW

Bases: Module

Convert tensor from (0:T, 1:H, 2:W, 3:C) to (3:C, 0:T, 1:H, 2:W)

Source code in zamba/pytorch/transforms.py

class ConvertTHWCtoCTHW(torch.nn.Module):
    """Convert tensor from (0:T, 1:H, 2:W, 3:C) to (3:C, 0:T, 1:H, 2:W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(3, 0, 1, 2)

ConvertTHWCtoTCHW

Bases: Module

Convert tensor from (T, H, W, C) to (T, C, H, W)

Source code in zamba/pytorch/transforms.py

class ConvertTHWCtoTCHW(torch.nn.Module):
    """Convert tensor from (T, H, W, C) to (T, C, H, W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(0, 3, 1, 2)

PackSlowFastPathways

Bases: Module

Creates the slow and fast pathway inputs for the slowfast model.

Source code in zamba/pytorch/transforms.py

class PackSlowFastPathways(torch.nn.Module):
    """Creates the slow and fast pathway inputs for the slowfast model."""

    def __init__(self, alpha: int = 4):
        super().__init__()
        self.alpha = alpha

    def forward(self, frames: torch.Tensor):
        fast_pathway = frames
        # Perform temporal sampling from the fast pathway.
        slow_pathway = torch.index_select(
            frames,
            1,
            torch.linspace(0, frames.shape[1] - 1, frames.shape[1] // self.alpha).long(),
        )
        frame_list = [slow_pathway, fast_pathway]
        return frame_list

PadDimensions

Bases: Module

Pads a tensor to ensure a fixed output dimension for a give axis.

Attributes:

Name	Type	Description
dimension_sizes		A tuple of int or None the same length as the number of dimensions in the input tensor. If int, pad that dimension to at least that size. If None, do not pad.

Source code in zamba/pytorch/transforms.py

class PadDimensions(torch.nn.Module):
    """Pads a tensor to ensure a fixed output dimension for a give axis.

    Attributes:
        dimension_sizes: A tuple of int or None the same length as the number of dimensions in the
            input tensor. If int, pad that dimension to at least that size. If None, do not pad.
    """

    def __init__(self, dimension_sizes: Tuple[Optional[int]]):
        super().__init__()
        self.dimension_sizes = dimension_sizes

    @staticmethod
    def compute_left_and_right_pad(original_size: int, padded_size: int) -> Tuple[int, int]:
        """Computes left and right pad size.

        Args:
            original_size (list, int): The original tensor size
            padded_size (list, int): The desired tensor size

        Returns:
           Tuple[int]: Pad size for right and left. For odd padding size, the right = left + 1
        """
        if original_size >= padded_size:
            return 0, 0
        pad = padded_size - original_size
        quotient, remainder = divmod(pad, 2)
        return quotient, quotient + remainder

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        padding = tuple(
            itertools.chain.from_iterable(
                (0, 0)
                if padded_size is None
                else self.compute_left_and_right_pad(original_size, padded_size)
                for original_size, padded_size in zip(vid.shape, self.dimension_sizes)
            )
        )
        return torch.nn.functional.pad(vid, padding[::-1])

compute_left_and_right_pad(original_size, padded_size) staticmethod

Computes left and right pad size.

Parameters:

Name	Type	Description	Default
original_size	(list, int)	The original tensor size	required
padded_size	(list, int)	The desired tensor size	required

Returns:

Type	Description
Tuple[int, int]	Tuple[int]: Pad size for right and left. For odd padding size, the right = left + 1

Source code in zamba/pytorch/transforms.py

@staticmethod
def compute_left_and_right_pad(original_size: int, padded_size: int) -> Tuple[int, int]:
    """Computes left and right pad size.

    Args:
        original_size (list, int): The original tensor size
        padded_size (list, int): The desired tensor size

    Returns:
       Tuple[int]: Pad size for right and left. For odd padding size, the right = left + 1
    """
    if original_size >= padded_size:
        return 0, 0
    pad = padded_size - original_size
    quotient, remainder = divmod(pad, 2)
    return quotient, quotient + remainder