Estimate image sharpness

sharpness.py

from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F

class Sharpness(nn.Module):
  def __init__(self):
    super().__init__()
    kernel = [[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]]
    kernel = Tensor(kernel)
    kernel.unsqueeze_(0)
    kernel.unsqueeze_(0)
    kernel = kernel.expand((3, 1, 3, 3))
    self.register_buffer('kernel', kernel)
    
  def forward(self, x):
    y = F.conv2d(x, self.kernel, groups=3)
    b, c, h, w = y.shape
    y = y.view(b, c, h * w)
    return y.std(dim=-1).pow(2).mean(dim=-1)