init

README.md

@@ -1,4 +1,5 @@
 ---
+python_version: 3.7
 title: U2net_portrait
 emoji: 🦀
 colorFrom: indigo

U-2-Net/utils/face_seg.py

@@ -0,0 +1,44 @@
+import os
+import cv2
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.platform import gfile
+
+
+curPath = os.path.abspath(os.path.dirname(__file__))
+
+
+class FaceSeg:
+    def __init__(self, model_path=os.path.join(curPath, 'seg_model_384.pb')):
+        config = tf.compat.v1.ConfigProto()
+        config.gpu_options.allow_growth = True
+        self._graph = tf.Graph()
+        self._sess = tf.compat.v1.Session(config=config, graph=self._graph)
+
+        self.pb_file_path = model_path
+        self._restore_from_pb()
+        self.input_op = self._sess.graph.get_tensor_by_name('input_1:0')
+        self.output_op = self._sess.graph.get_tensor_by_name('sigmoid/Sigmoid:0')
+
+    def _restore_from_pb(self):
+        with self._sess.as_default():
+            with self._graph.as_default():
+                with gfile.FastGFile(self.pb_file_path, 'rb') as f:
+                    graph_def = tf.compat.v1.GraphDef()
+                    graph_def.ParseFromString(f.read())
+                    tf.import_graph_def(graph_def, name='')
+
+    def input_transform(self, image):
+        image = cv2.resize(image, (384, 384), interpolation=cv2.INTER_AREA)
+        image_input = (image / 255.)[np.newaxis, :, :, :]
+        return image_input
+
+    def output_transform(self, output, shape):
+        output = cv2.resize(output, (shape[1], shape[0]))
+        image_output = (output * 255).astype(np.uint8)
+        return image_output
+
+    def get_mask(self, image):
+        image_input = self.input_transform(image)
+        output = self._sess.run(self.output_op, feed_dict={self.input_op: image_input})[0]
+        return self.output_transform(output, shape=image.shape[:2])

U-2-Net/utils/seg_model_384.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66a04bc2032b54013d2ae994b34d22518144276f1cbdd2d8cbb1a4a28f50285f
+size 32477258

app.py

@@ -27,8 +27,9 @@ from data_loader import SalObjDataset
 from model import U2NET # full size version 173.6 MB
 from model import U2NETP # small version u2net 4.7 MB
 
-from modnet import ModNet
 import huggingface_hub
+from utils.face_seg import FaceSeg
+import cv2
 
 # normalize the predicted SOD probability map
 def normPRED(d):
@@ -59,17 +60,11 @@ def save_output(image_name,pred,d_dir):
     imo.save(d_dir+'/'+imidx+'.png')
     return d_dir+'/'+imidx+'.png'
 
-
-
-modnet_path = huggingface_hub.hf_hub_download('hylee/apdrawing_model',
-                                                      'modnet.onnx',
-                                                      force_filename='modnet.onnx')
-modnet = ModNet(modnet_path)
+segment = FaceSeg()
 
  # --------- 1. get image path and name ---------
 model_name='u2net_portrait'#u2netp
 
-
 image_dir = 'portrait_im'
 prediction_dir = 'portrait_results'
 if(not os.path.exists(prediction_dir)):
@@ -90,9 +85,18 @@ net.eval()
 
 
 def process(im):
-    #image = modnet.segment(im.name)
-    #im_path = os.path.abspath(os.path.basename(im.name))
-    #Image.fromarray(np.uint8(image)).save(im_path)
+    image = cv2.imread(im.name)
+    matte = self.segment.get_mask(face)
+
+    if len(image.shape) == 2:
+        image = image[:, :, None]
+    if image.shape[2] == 1:
+        image = np.repeat(image, 3, axis=2)
+    elif image.shape[2] == 4:
+        image = image[:, :, 0:3]
+    matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2) / 255
+    foreground = image * matte + np.full(image.shape, 255) * (1 - matte)
+    cv2.imwrite(im.name, foreground)
 
     img_name_list = [im.name]
     print("Number of images: ", len(img_name_list))
@@ -135,7 +139,7 @@ def process(im):
 
     print(results)
 
-    return Image.open(results[0])
+    return Image.open(results[0]), Image.open(im.name)
         
 title = "U-2-Net"
 description = "Gradio demo for U-2-Net, https://github.com/xuebinqin/U-2-Net"
@@ -145,7 +149,7 @@ gr.Interface(
     process, 
     [gr.inputs.Image(type="file", label="Input")
 ], 
-    [gr.outputs.Image(type="pil", label="Output")],
+    [gr.outputs.Image(type="pil", label="Output"), gr.outputs.Image(type="pil", label="Output")],
     title=title,
     description=description,
     article=article,

requirements.txt

@@ -4,5 +4,4 @@ torch
 torchvision
 pillow
 opencv-python-headless
-onnx==1.8.1
-onnxruntime==1.6.0
+tensorflow-gpu==1.14.0