* Return a sharpened version of the image, using an unsharp mask

.gitattributes

@@ -25,5 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-u2net_bce_itr_18000_train_3.891670_tar_0.553700_512x_460x.jit.pt filter=lfs diff=lfs merge=lfs -text
+u2net_bce_itr_25000_train_3.856416_tar_0.547567-400x_360x.jit.pt filter=lfs diff=lfs merge=lfs -text
 shape_predictor_5_face_landmarks.dat filter=lfs diff=lfs merge=lfs -text

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_18000_train_3.891670_tar_0.553700_512x_460x.jit.pt')
+model = torch.jit.load('u2net_bce_itr_25000_train_3.856416_tar_0.547567-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)) 
@@ -45,11 +66,11 @@ def image_as_array(image_in):
     image_out = np.expand_dims(tmpImg, 0)
     return image_out
 
-def find_aligned_face(image_in, size=512):   
+def find_aligned_face(image_in, size=400):   
     aligned_image, n_faces, quad = face_detection.align(image_in, face_index=0, output_size=size)
     return aligned_image, n_faces, quad
 
-def align_first_face(image_in, size=512):
+def align_first_face(image_in, size=400):
     aligned_image, n_faces, quad = find_aligned_face(image_in,size=size)
     if n_faces == 0:
         try:
@@ -82,14 +103,17 @@ def face2doll(
     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)         
+        doll_np_image = array_to_np(results[1].detach().numpy())
+        doll_image = unsharp_mask(doll_np_image)
+        doll_image = Image.fromarray(doll_image)
+        
+        output = img_concat_h(array_to_image(aligned_img), doll_image)         
         del results
 
     return output
 
 def inference(img):
-    out = face2doll(img, 512)
+    out = face2doll(img, 400)
     return out
       
   

u2net_bce_itr_18000_train_3.891670_tar_0.553700_512x_460x.jit.pt → u2net_bce_itr_25000_train_3.856416_tar_0.547567-400x_360x.jit.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a4759db0de410078514769bdae185cf80b095e8a5b49aef31987ad5fa6e66a30
-size 177290264
+oid sha256:9b8d7427e71fd7cb303ffffaef7a42c7c7f95aa77f846fa4cccb4130fcfbbf74
+size 177193974