init

app.py

@@ -112,6 +112,7 @@ def compress_UUID():
 
 def run(
         image,
+        mask,
         model,
         opt,
         detector,
@@ -124,8 +125,10 @@ def run(
     os.makedirs(opt.results_dir, exist_ok=True)
 
     opt.lm_dir = os.path.join(dataroot, 'landmark/')
-    opt.bg_dir = os.path.join(dataroot, 'mask/')
+    opt.cmask_dir = os.path.join(dataroot, 'mask/')
+    opt.bg_dir = os.path.join(dataroot, 'mask/bg')
     os.makedirs(opt.lm_dir, exist_ok=True)
+    os.makedirs(opt.cmask_dir, exist_ok=True)
     os.makedirs(opt.bg_dir, exist_ok=True)
 
     shutil.copy(image.name, opt.dataroot)
@@ -133,6 +136,8 @@ def run(
     fullname = os.path.basename(image.name)
     name = fullname.split(".")[0]
 
+    shutil.copyfile(mask.name, os.path.join(opt.bg_dir, name+'.png'))
+
     imgfile = os.path.join(opt.dataroot, fullname)
     lmfile = os.path.join(opt.lm_dir, name+'.txt')
     # 预处理数据
@@ -140,30 +145,29 @@ def run(
 
     imgs = []
     for part in ['eyel', 'eyer', 'nose', 'mouth']:
-        savepath = os.path.join(opt.bg_dir + part, name+'.png')
+        savepath = os.path.join(opt.cmask_dir + part, name+'.png')
         get_partmask(imgfile, part, lmfile, savepath)
-        imgs.append(savepath)
-
-    # data_loader = CreateDataLoader(opt)
-    # dataset = data_loader.load_data()
-    #
-    # imgs = [image.name]
-    # # test
-    # # model.eval()
-    # for i, data in enumerate(dataset):
-    #     if i >= opt.how_many:  # test code only supports batch_size = 1, how_many means how many test images to run
-    #         break
-    #     model.set_input(data)
-    #     model.test()
-    #     visuals = model.get_current_visuals()  # in test the loadSize is set to the same as fineSize
-    #     img_path = model.get_image_paths()
-    #     # if i % 5 == 0:
-    #     #    print('processing (%04d)-th image... %s' % (i, img_path))
-    #     imgs = save_images2(opt.results_dir, visuals, img_path, aspect_ratio=opt.aspect_ratio, width=opt.display_winsize)
-    #
+        #imgs.append(savepath)
+
+    data_loader = CreateDataLoader(opt)
+    dataset = data_loader.load_data()
+
+    # test
+    # model.eval()
+    for i, data in enumerate(dataset):
+        if i >= opt.how_many:  # test code only supports batch_size = 1, how_many means how many test images to run
+            break
+        model.set_input(data)
+        model.test()
+        visuals = model.get_current_visuals()  # in test the loadSize is set to the same as fineSize
+        img_path = model.get_image_paths()
+        # if i % 5 == 0:
+        #    print('processing (%04d)-th image... %s' % (i, img_path))
+        imgs = save_images2(opt.results_dir, visuals, img_path, aspect_ratio=opt.aspect_ratio, width=opt.display_winsize)
+
     print(imgs)
 
-    return PIL.Image.open(imgs[0]),PIL.Image.open(imgs[1]),PIL.Image.open(imgs[2]),PIL.Image.open(imgs[3])
+    return PIL.Image.open(imgs[0]),PIL.Image.open(imgs[1])
 
 
 def main():
@@ -212,6 +216,7 @@ def main():
         func,
         [
             gr.inputs.Image(type='file', label='Input Image'),
+            gr.inputs.Image(type='file', label='Input Mask'),
         ],
         [
             gr.outputs.Image(