Sharper output with unsharp mask

app.py

@@ -6,12 +6,27 @@ import face_detection
 import PIL
 from PIL import Image, ImageOps, ImageFile
 import numpy as np
-
+import cv2 as cv
 import torch
+
 torch.set_grad_enabled(False)
 model = torch.jit.load('u2net_bce_itr_16000_train_3.835149_tar_0.542587-400x_360x.jit.pt')
 model.eval()
 
+# https://en.wikipedia.org/wiki/Unsharp_masking
+# https://stackoverflow.com/a/55590133/1495606
+def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0):
+    """Return a sharpened version of the image, using an unsharp mask."""
+    blurred = cv.GaussianBlur(image, kernel_size, sigma)
+    sharpened = float(amount + 1) * image - float(amount) * blurred
+    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
+    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
+    sharpened = sharpened.round().astype(np.uint8)
+    if threshold > 0:
+        low_contrast_mask = np.absolute(image - blurred) < threshold
+        np.copyto(sharpened, image, where=low_contrast_mask)
+    return sharpened
+
 def normPRED(d):
     ma = np.max(d)
     mi = np.min(d)
@@ -20,6 +35,12 @@ def normPRED(d):
 
     return dn
 
+def array_to_np(array_in):
+    array_in = normPRED(array_in)
+    array_in = np.squeeze(255.0*(array_in)) 
+    array_in = np.transpose(array_in, (1, 2, 0))
+    return array_in
+
 def array_to_image(array_in):    
     array_in = normPRED(array_in)
     array_in = np.squeeze(255.0*(array_in)) 
@@ -82,8 +103,11 @@ def face2hero(
     else:
         input = torch.Tensor(aligned_img)
         results = model(input)
-        d2 = array_to_image(results[1].detach().numpy())
-        output = img_concat_h(array_to_image(aligned_img), d2)         
+        hero_np_image = array_to_np(results[1].detach().numpy())
+        hero_image = unsharp_mask(hero_np_image)
+        hero_image = Image.fromarray(hero_image)
+        
+        output = img_concat_h(array_to_image(aligned_img), hero_image)         
         del results
 
     return output