init

app.py

@@ -29,6 +29,7 @@ from util import html
 
 import ntpath
 from util import util
+from modnet import ModNet
 
 ORIGINAL_REPO_URL = 'https://github.com/yiranran/APDrawingGAN2'
 TITLE = 'yiranran/APDrawingGAN2'
@@ -64,17 +65,15 @@ def load_checkpoint():
                                                       force_filename='checkpoints.zip')
     print(checkpoint_path)
     shutil.unpack_archive(checkpoint_path, extract_dir=dir)
-
     print(os.listdir(dir + '/checkpoints'))
-
-    dataset_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                      'dataset.zip',
-                                                      force_filename='dataset.zip')
-    print(checkpoint_path)
-    shutil.unpack_archive(dataset_path)
-
     return dir + '/checkpoints'
 
+def load_modnet_model():
+    modnet_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                      'modnet.onnx',
+                                                      force_filename='modnet.onnx')
+    return modnet_path
+
 
 # save image to the disk
 def save_images2(image_dir, visuals, image_path, aspect_ratio=1.0, width=256):
@@ -124,6 +123,7 @@ def run(
         opt,
         detector,
         predictor,
+        modnet : ModNet,
 ) -> tuple[PIL.Image.Image,PIL.Image.Image]:
 
     dataroot = 'images/' + compress_UUID()
@@ -144,9 +144,9 @@ def run(
     fullname = os.path.basename(image.name)
     name = fullname.split(".")[0]
 
-    bg = cv2.cvtColor(cv2.imread(mask.name), cv2.COLOR_BGR2GRAY)
-    cv2.imwrite(os.path.join(opt.bg_dir, name+'.png'), bg)
-    #shutil.copyfile(mask.name, os.path.join(opt.bg_dir, name+'.png'))
+    #bg = cv2.cvtColor(cv2.imread(mask.name), cv2.COLOR_BGR2GRAY)
+    #cv2.imwrite(os.path.join(opt.bg_dir, name+'.png'), bg)
+    modnet.segment(image.name, os.path.join(opt.bg_dir, name+'.png'))
 
     imgfile = os.path.join(opt.dataroot, fullname)
     lmfile5 = os.path.join(opt.lm_dir, name+'.txt')
@@ -155,10 +155,10 @@ def run(
     get_68lm(imgfile, lmfile5, lmfile68, detector, predictor)
 
     imgs = []
-    # for part in ['eyel', 'eyer', 'nose', 'mouth']:
-    #     savepath = os.path.join(opt.cmask_dir + part, name+'.png')
-    #     get_partmask(imgfile, part, lmfile68, savepath)
-    #     #imgs.append(savepath)
+    for part in ['eyel', 'eyer', 'nose', 'mouth']:
+        savepath = os.path.join(opt.cmask_dir + part, name+'.png')
+        get_partmask(imgfile, part, lmfile68, savepath)
+        #imgs.append(savepath)
 
     data_loader = CreateDataLoader(opt)
     dataset = data_loader.load_data()
@@ -201,13 +201,17 @@ def main():
     model = create_model(opt)
     model.setup(opt)
 
+    modnet_path = load_modnet_model();
+    modnet = ModNet(modnet_path)
+
+
     '''
     预处理数据
     '''
     detector = dlib.get_frontal_face_detector()
     predictor = dlib.shape_predictor(checkpoint_dir + '/shape_predictor_68_face_landmarks.dat')
 
-    func = functools.partial(run, model=model, opt=opt, detector=detector, predictor=predictor)
+    func = functools.partial(run, model=model, opt=opt, detector=detector, predictor=predictor, modnet=modnet)
     func = functools.update_wrapper(func, run)
 
     gr.Interface(

requirements.txt

@@ -7,4 +7,6 @@ numpy==1.16.4
 pillow<7.0.0
 opencv-python==4.1.0.25
 dlib==19.18.0
-shapely==1.7.0
+shapely==1.7.0
+onnx==1.8.1
+onnxruntime==1.6.0