def smooth_geom()

in utils/helpers.py [0:0]

• 17 lines of code
• 2 McCabe index (conditional complexity)

def smooth_geom(geom, mask: th.Tensor = None, filter_size: int = 9, sigma: float = 2.0):
    """
    :param geom: T x V x 3 tensor containing a temporal sequence of length T with V vertices in each frame
    :param mask: V-dimensional Tensor containing a mask with vertices to be smoothed
    :param filter_size: size of the Gaussian filter
    :param sigma: standard deviation of the Gaussian filter
    :return: T x V x 3 tensor containing smoothed geometry (i.e., smoothed in the area indicated by the mask)
    """
    assert filter_size % 2 == 1, f"filter size must be odd but is {filter_size}"
    # Gaussian smoothing (low-pass filtering)
    fltr = np.arange(-(filter_size // 2), filter_size // 2 + 1)
    fltr = np.exp(-0.5 * fltr ** 2 / sigma ** 2)
    fltr = th.Tensor(fltr) / np.sum(fltr)
    # apply fltr
    fltr = fltr.view(1, 1, -1).to(device=geom.device)
    T, V = geom.shape[0], geom.shape[1]
    g = th.nn.functional.pad(
        geom.permute(1, 2, 0).view(V * 3, 1, T),
        pad=[filter_size // 2, filter_size // 2], mode='replicate'
    )
    g = th.nn.functional.conv1d(g, fltr).view(V, 3, T)
    smoothed = g.permute(2, 0, 1).contiguous()
    # blend smoothed signal with original signal
    if mask is None:
        return smoothed
    else:
        return smoothed * mask[None, :, None] + geom * (-mask[None, :, None] + 1)