|
|
|
|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
import torch |
|
|
from torch.nn import functional as F |
|
|
|
|
|
import os, sys |
|
|
root_path = os.path.abspath('.') |
|
|
sys.path.append(root_path) |
|
|
from degradation.ESR.utils import filter2D, np2tensor, tensor2np |
|
|
|
|
|
|
|
|
def usm_sharp_func(img, weight=0.5, radius=50, threshold=10): |
|
|
"""USM sharpening. |
|
|
|
|
|
Input image: I; Blurry image: B. |
|
|
1. sharp = I + weight * (I - B) |
|
|
2. Mask = 1 if abs(I - B) > threshold, else: 0 |
|
|
3. Blur mask: |
|
|
4. Out = Mask * sharp + (1 - Mask) * I |
|
|
|
|
|
|
|
|
Args: |
|
|
img (Numpy array): Input image, HWC, BGR; float32, [0, 1]. |
|
|
weight (float): Sharp weight. Default: 1. |
|
|
radius (float): Kernel size of Gaussian blur. Default: 50. |
|
|
threshold (int): |
|
|
""" |
|
|
if radius % 2 == 0: |
|
|
radius += 1 |
|
|
blur = cv2.GaussianBlur(img, (radius, radius), 0) |
|
|
residual = img - blur |
|
|
mask = np.abs(residual) * 255 > threshold |
|
|
mask = mask.astype('float32') |
|
|
soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0) |
|
|
|
|
|
sharp = img + weight * residual |
|
|
sharp = np.clip(sharp, 0, 1) |
|
|
return soft_mask * sharp + (1 - soft_mask) * img |
|
|
|
|
|
|
|
|
|
|
|
class USMSharp(torch.nn.Module): |
|
|
|
|
|
def __init__(self, type, radius=50, sigma=0): |
|
|
super(USMSharp, self).__init__() |
|
|
if radius % 2 == 0: |
|
|
radius += 1 |
|
|
self.radius = radius |
|
|
kernel = cv2.getGaussianKernel(radius, sigma) |
|
|
kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0).cuda() |
|
|
self.register_buffer('kernel', kernel) |
|
|
|
|
|
self.type = type |
|
|
|
|
|
|
|
|
def forward(self, img, weight=0.5, threshold=10, store=False): |
|
|
|
|
|
if self.type == "cv2": |
|
|
|
|
|
img = np2tensor(img) |
|
|
|
|
|
blur = filter2D(img, self.kernel.cuda()) |
|
|
if store: |
|
|
cv2.imwrite("blur.png", tensor2np(blur)) |
|
|
|
|
|
residual = img - blur |
|
|
if store: |
|
|
cv2.imwrite("residual.png", tensor2np(residual)) |
|
|
|
|
|
mask = torch.abs(residual) * 255 > threshold |
|
|
if store: |
|
|
cv2.imwrite("mask.png", tensor2np(mask)) |
|
|
|
|
|
|
|
|
mask = mask.float() |
|
|
soft_mask = filter2D(mask, self.kernel.cuda()) |
|
|
if store: |
|
|
cv2.imwrite("soft_mask.png", tensor2np(soft_mask)) |
|
|
|
|
|
sharp = img + weight * residual |
|
|
sharp = torch.clip(sharp, 0, 1) |
|
|
if store: |
|
|
cv2.imwrite("sharp.png", tensor2np(sharp)) |
|
|
|
|
|
output = soft_mask * sharp + (1 - soft_mask) * img |
|
|
if self.type == "cv2": |
|
|
output = tensor2np(output) |
|
|
|
|
|
return output |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
|
|
usm_sharper = USMSharp(type="cv2") |
|
|
img = cv2.imread("sample3.png") |
|
|
print(img.shape) |
|
|
sharp_output = usm_sharper(img, store=False, threshold=10) |
|
|
cv2.imwrite(os.path.join("output.png"), sharp_output) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
