update the exmaples

.gitignore

@@ -0,0 +1,2 @@
+**/*.*.pyc
+**/*.pyc

app.py

@@ -37,6 +37,7 @@ def resize(img, size):
         newh = int(h / w * size)
 
     resized_img = cv2.resize(img, (neww, newh), interpolation=cv2.INTER_AREA)
+
     if len(img.shape) != len(resized_img.shape):
         resized_img = resized_img[..., none]
 
@@ -62,7 +63,7 @@ def padding_mask(rgba_input: np.array):
     :returns: H x W x 4 padded RGBAD
 
     """
-    padding = 50
+    padding = 40
     padding_size = 256 - padding * 2
 
     h, w = rgba_input.shape[:2]
@@ -74,6 +75,7 @@ def padding_mask(rgba_input: np.array):
     h_min, h_max = hh.min(), hh.max()
     w_min, w_max = ww.min(), ww.max()
 
+    
     # if the area already has enough padding
     if h_max - h_min < padding_size and w_max - w_min < padding_size:
         return rgba_input
@@ -83,8 +85,12 @@ def padding_mask(rgba_input: np.array):
 
     padded_rgba  = resize(rgba_input, padding_size)
     new_h, new_w = padded_rgba.shape[:2]
+    
+    padding_h = (256 - new_h) // 2
+    padding_w = (256 - new_w) // 2
 
-    padding_output[padding:padding+new_h, padding:padding+new_w, :] = padded_rgba
+    padding_output[padding_h:padding_h+new_h, padding_w:padding_w+new_w, :] = padded_rgba
+    padding_output = np.clip(padding_output, 0.0, 1.0)
 
     return padding_output
 
@@ -108,6 +114,8 @@ def render_btn_fn(mask, ibl):
     mask = mask / 255.0
     ibl = ibl/ 255.0
 
+    mask = np.clip(mask, 0.0, 1.0)
+
     # smoothing ibl
     ibl = cv2.GaussianBlur(ibl, (11, 11), 0)
 
@@ -157,13 +165,16 @@ def gamma_change(x):
     ret, shadow = shadow_composite(cur_rgba, cur_shadow, cur_intensity, cur_gamma)    
     return ret, shadow
 
+def update_input(mask):
+    return mask
+
 
 ibl_h = 128
 ibl_w = ibl_h * 2
 
 with gr.Blocks() as demo:
     with gr.Row():
-        mask_input = gr.Image(shape=(256, 256), image_mode="RGBA", label="Mask")
+        mask_input = gr.Image(shape=None, width=256, height=256,image_mode="RGBA", label="RGBA")
         ibl_input = gr.Sketchpad(shape=(ibl_w, ibl_h), image_mode="L", label="IBL", tool='sketch', invert_colors=True)
         output = gr.Image(shape=(256, 256), height=256, width=256, image_mode="RGB", label="Output")
         shadow_output = gr.Image(shape=(256, 256), height=256, width=256, image_mode="L", label="Shadow Layer")
@@ -173,6 +184,14 @@ with gr.Blocks() as demo:
         gamma_slider = gr.Slider(1.0, 4.0, value=DEFAULT_GAMMA, step=0.1, label="Gamma", info="Gamma correction for shadow") 
         render_btn = gr.Button(label="Render")
 
+    with gr.Row():
+        gr.Examples(
+            examples=[['imgs/woman.png'],['imgs/man.png'], ['imgs/plant1.png'], ['imgs/human2.png'], ['imgs/cloud.png']],
+            fn=update_input, 
+            inputs=[mask_input],
+            outputs=mask_input
+        )
+
     render_btn.click(render_btn_fn, inputs=[mask_input, ibl_input], outputs=[output, shadow_output])
     intensity_slider.release(intensity_change, inputs=[intensity_slider], outputs=[output, shadow_output])
     gamma_slider.release(gamma_change, inputs=[gamma_slider], outputs=[output, shadow_output])

imgs/convert_rgb_alpha.py

@@ -0,0 +1,14 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+rgb_file = 'fg-1-rgb.png'
+alpha_file = 'fg-1-alpha.png'
+output_file = 'fg-1-rgba.png'
+
+rgb = plt.imread(rgb_file)
+alpha = plt.imread(alpha_file)
+
+print(rgb.shape, alpha.shape)
+
+rgba = np.concatenate([rgb[..., :3], alpha[..., 0:1]], axis=2)
+plt.imsave(output_file, rgba)

imgs/test.py

@@ -0,0 +1,14 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+rgb = 'woman.png'
+mask = 'woman_mask.png'
+ofile = 'test1.png'
+
+rgb = plt.imread(rgb)
+mask = plt.imread(mask)
+
+output = np.concatenate([rgb[..., :3], mask[..., :1]], axis=2)
+plt.imsave(ofile, output)
+

models/Loss/Loss.py

@@ -10,7 +10,7 @@ import cv2
 # from vgg19_loss import VGG19Loss
 # import pytorch_ssim
 
-from .vgg19_loss import VGG19Loss
+# from .vgg19_loss import VGG19Loss
 from . import pytorch_ssim
 from abc import ABC, abstractmethod
 from collections import OrderedDict
@@ -96,15 +96,15 @@ class hierarchical_ssim_loss(abs_loss):
         return total_loss/b
 
 
-class vgg_loss(abs_loss):
-    def __init__(self):
-        self.vgg19_ = VGG19Loss()
+# class vgg_loss(abs_loss):
+#     def __init__(self):
+#         self.vgg19_ = VGG19Loss()
 
 
-    def loss(self, gt_img, pred_img):
-        b, c, h, w = gt_img.shape
-        v = self.vgg19_(gt_img, pred_img, pred_img.device)
-        return v/b
+#     def loss(self, gt_img, pred_img):
+#         b, c, h, w = gt_img.shape
+#         v = self.vgg19_(gt_img, pred_img, pred_img.device)
+#         return v/b
 
 
 class grad_loss(abs_loss):

models/SSN.py

@@ -97,7 +97,8 @@ class SSN(abs_model):
             assert len(x[k].shape) == 2, '{} should be 2D tensor'.format(k)
 
 
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        device = torch.device('cpu')
 
         mask = torch.tensor(x['mask'])[None, None, ...].float().to(device)
         ibl  = torch.tensor(x['ibl'])[None, None, ...].float().to(device)