40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
|
|
from mmselfsup.models.target_generators import HOGGenerator
|
|
from mmselfsup.registry import MODELS
|
|
|
|
|
|
@MODELS.register_module()
|
|
class HOGGenerator3d(HOGGenerator):
|
|
"""Generate HOG feature for videos.
|
|
|
|
This module is used in MaskFeat to generate HOG feature.
|
|
Here is the link of `HOG wikipedia
|
|
<https://en.wikipedia.org/wiki/Histogram_of_oriented_gradients>`_.
|
|
|
|
Args:
|
|
nbins (int): Number of bin. Defaults to 9.
|
|
pool (float): Number of cell. Defaults to 8.
|
|
gaussian_window (int): Size of gaussian kernel. Defaults to 16.
|
|
"""
|
|
|
|
def __init__(self,
|
|
nbins: int = 9,
|
|
pool: int = 8,
|
|
gaussian_window: int = 16) -> None:
|
|
super().__init__(
|
|
nbins=nbins, pool=pool, gaussian_window=gaussian_window)
|
|
|
|
def _reshape(self, hog_feat: torch.Tensor) -> torch.Tensor:
|
|
"""Reshape HOG Features for output."""
|
|
hog_feat = hog_feat.flatten(1, 2)
|
|
self.unfold_size = hog_feat.shape[-1] // 14
|
|
hog_feat = hog_feat.permute(0, 2, 3, 1)
|
|
hog_feat = hog_feat.unfold(1, self.unfold_size,
|
|
self.unfold_size).unfold(
|
|
2, self.unfold_size, self.unfold_size)
|
|
hog_feat = hog_feat.flatten(3).view(self.B, self.T, 14, 14, -1)
|
|
hog_feat = hog_feat.flatten(1, 3) # B N C
|
|
return hog_feat
|