mmselfsup.models.target_generators.hog_generator 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import math
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from mmengine.model import BaseModule
from mmselfsup.registry import MODELS
[文档]@MODELS.register_module()
class HOGGenerator(BaseModule):
"""Generate HOG feature for images.
This module is used in MaskFeat to generate HOG feature. The code is
modified from file `slowfast/models/operators.py
<https://github.com/facebookresearch/SlowFast/blob/main/slowfast/models/operators.py>`_.
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__()
self.nbins = nbins
self.pool = pool
self.pi = math.pi
weight_x = torch.FloatTensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
weight_x = weight_x.view(1, 1, 3, 3).repeat(3, 1, 1, 1).contiguous()
weight_y = weight_x.transpose(2, 3).contiguous()
self.register_buffer('weight_x', weight_x)
self.register_buffer('weight_y', weight_y)
self.gaussian_window = gaussian_window
if gaussian_window:
gaussian_kernel = self.get_gaussian_kernel(gaussian_window,
gaussian_window // 2)
self.register_buffer('gaussian_kernel', gaussian_kernel)
[文档] def get_gaussian_kernel(self, kernlen: int, std: int) -> torch.Tensor:
"""Returns a 2D Gaussian kernel array."""
def _gaussian_fn(kernlen: int, std: int) -> torch.Tensor:
n = torch.arange(0, kernlen).float()
n -= n.mean()
n /= std
w = torch.exp(-0.5 * n**2)
return w
kernel_1d = _gaussian_fn(kernlen, std)
kernel_2d = kernel_1d[:, None] * kernel_1d[None, :]
return kernel_2d / kernel_2d.sum()
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(1, 2).flatten(2)
return hog_feat
[文档] @torch.no_grad()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Generate hog feature for each batch images.
Args:
x (torch.Tensor): Input images of shape (N, 3, H, W).
Returns:
torch.Tensor: Hog features.
"""
# input is RGB image with shape [B 3 H W]
self.h, self.w = x.size(-2), x.size(-1)
x = F.pad(x, pad=(1, 1, 1, 1), mode='reflect')
gx_rgb = F.conv2d(
x, self.weight_x, bias=None, stride=1, padding=0, groups=3)
gy_rgb = F.conv2d(
x, self.weight_y, bias=None, stride=1, padding=0, groups=3)
norm_rgb = torch.stack([gx_rgb, gy_rgb], dim=-1).norm(dim=-1)
phase = torch.atan2(gx_rgb, gy_rgb)
phase = phase / self.pi * self.nbins # [-9, 9]
b, c, h, w = norm_rgb.shape
out = torch.zeros((b, c, self.nbins, h, w),
dtype=torch.float,
device=x.device)
phase = phase.view(b, c, 1, h, w)
norm_rgb = norm_rgb.view(b, c, 1, h, w)
if self.gaussian_window:
if h != self.gaussian_window:
assert h % self.gaussian_window == 0, 'h {} gw {}'.format(
h, self.gaussian_window)
repeat_rate = h // self.gaussian_window
temp_gaussian_kernel = self.gaussian_kernel.repeat(
[repeat_rate, repeat_rate])
else:
temp_gaussian_kernel = self.gaussian_kernel
norm_rgb *= temp_gaussian_kernel
out.scatter_add_(2, phase.floor().long() % self.nbins, norm_rgb)
out = out.unfold(3, self.pool, self.pool)
out = out.unfold(4, self.pool, self.pool)
out = out.sum(dim=[-1, -2])
self.out = F.normalize(out, p=2, dim=2)
return self._reshape(self.out)
[文档] def generate_hog_image(self, hog_out: torch.Tensor) -> np.ndarray:
"""Generate HOG image according to HOG features."""
assert hog_out.size(0) == 1 and hog_out.size(1) == 3, \
'Check the input batch size and the channcel number, only support'\
'"batch_size = 1".'
hog_image = np.zeros([self.h, self.w])
cell_gradient = np.array(hog_out.mean(dim=1).squeeze().detach().cpu())
cell_width = self.pool / 2
max_mag = np.array(cell_gradient).max()
angle_gap = 360 / self.nbins
for x in range(cell_gradient.shape[1]):
for y in range(cell_gradient.shape[2]):
cell_grad = cell_gradient[:, x, y]
cell_grad /= max_mag
angle = 0
for magnitude in cell_grad:
angle_radian = math.radians(angle)
x1 = int(x * self.pool +
magnitude * cell_width * math.cos(angle_radian))
y1 = int(y * self.pool +
magnitude * cell_width * math.sin(angle_radian))
x2 = int(x * self.pool -
magnitude * cell_width * math.cos(angle_radian))
y2 = int(y * self.pool -
magnitude * cell_width * math.sin(angle_radian))
magnitude = 0 if magnitude < 0 else magnitude
cv2.line(hog_image, (y1, x1), (y2, x2),
int(255 * math.sqrt(magnitude)))
angle += angle_gap
return hog_image