From a27b349938ed23c6223962afce302dd9ed2dfe56 Mon Sep 17 00:00:00 2001
From: EllangoK <karun.ellango7@gmail.com>
Date: Fri, 31 Mar 2023 12:22:23 -0400
Subject: [PATCH] pure pytorch sharpen node

---
 sharpen.py | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/sharpen.py b/sharpen.py
index a8df732..824094e 100644
--- a/sharpen.py
+++ b/sharpen.py
@@ -1,6 +1,5 @@
-import cv2
-import numpy as np
 import torch
+import torch.nn.functional as F
 
 
 class Sharpen:
@@ -33,21 +32,21 @@ class Sharpen:
     CATEGORY = "postprocessing"
 
     def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float):
-        batch_size, height, width, _ = image.shape
+        batch_size, height, width, channels = image.shape
         result = torch.zeros_like(image)
 
+        kernel = torch.ones((channels, 1, kernel_size, kernel_size), dtype=torch.float32) * -1
+        center = kernel_size // 2
+        kernel[:, 0, center, center] = kernel_size**2
+        kernel *= alpha
+
         for b in range(batch_size):
-            tensor_image = image[b].numpy()
+            tensor_image = image[b].permute(2, 0, 1).unsqueeze(0)
 
-            kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1
-            center = kernel_size // 2
-            kernel[center, center] = kernel_size**2
-            kernel *= alpha
+            sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)
+            sharpened = sharpened.squeeze(0).permute(1, 2, 0)
 
-            sharpened = cv2.filter2D(tensor_image, -1, kernel)
-
-            tensor = torch.from_numpy(sharpened).unsqueeze(0)
-            tensor = torch.clamp(tensor, 0, 1)
+            tensor = torch.clamp(sharpened, 0, 1)
             result[b] = tensor
 
         return (result,)