Module facetorch.transforms

class SquarePad(torch.nn.Module):
    """SquarePad is a transform that pads the image to a square shape."""

    def __init__(self) -> None:
        """It is initialized as a torch.nn.Module."""
        super().__init__()

    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
        """Pads a tensor to a square.

        Args:
            tensor (torch.Tensor): tensor to pad.

        Returns:
            torch.Tensor: Padded tensor.
        """
        height, width = tensor.shape[-2:]
        img_size = [width, height]

        max_wh = max(img_size)
        p_left, p_top = [(max_wh - s) // 2 for s in img_size]
        p_right, p_bottom = [
            max_wh - (s + pad) for s, pad in zip(img_size, [p_left, p_top])
        ]
        padding = (p_left, p_top, p_right, p_bottom)
        tensor_padded = torchvision.transforms.functional.pad(
            tensor, padding, 0, "constant"
        )
        return tensor_padded

    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
        """Pads a tensor to a square.

        Args:
            tensor (torch.Tensor): tensor to pad.

        Returns:
            torch.Tensor: Padded tensor.

        """
        return self.__call__(tensor)