init

app.py

@@ -22,13 +22,13 @@ import shutil
 from options.test_options import TestOptions
 from data import CreateDataLoader
 from models import create_model
-
+import dlib
+import preprocess.get_partmask
 from util import html
 
 import ntpath
 from util import util
 
-
 ORIGINAL_REPO_URL = 'https://github.com/yiranran/APDrawingGAN2'
 TITLE = 'yiranran/APDrawingGAN2'
 DESCRIPTION = f"""This is a demo for {ORIGINAL_REPO_URL}.
@@ -38,9 +38,9 @@ ARTICLE = """
 
 """
 
-
 MODEL_REPO = 'hylee/apdrawing_model'
 
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument('--device', type=str, default='cpu')
@@ -59,14 +59,15 @@ def parse_args() -> argparse.Namespace:
 def load_checkpoint():
     dir = 'checkpoint'
     checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                'checkpoints.zip',
-                                                force_filename='checkpoints.zip')
+                                                      'checkpoints.zip',
+                                                      force_filename='checkpoints.zip')
     print(checkpoint_path)
     shutil.unpack_archive(checkpoint_path, extract_dir=dir)
 
-    print(os.listdir(dir+'/checkpoints'))
+    print(os.listdir(dir + '/checkpoints'))
+
+    return dir + '/checkpoints'
 
-    return dir+'/checkpoints'
 
 # save image to the disk
 def save_images2(image_dir, visuals, image_path, aspect_ratio=1.0, width=256):
@@ -76,7 +77,7 @@ def save_images2(image_dir, visuals, image_path, aspect_ratio=1.0, width=256):
     imgs = []
 
     for label, im_data in visuals.items():
-        im = util.tensor2im(im_data)#tensor to numpy array [-1,1]->[0,1]->[0,255]
+        im = util.tensor2im(im_data)  # tensor to numpy array [-1,1]->[0,1]->[0,255]
         image_name = '%s_%s.png' % (name, label)
         save_path = os.path.join(image_dir, image_name)
         h, w, _ = im.shape
@@ -91,6 +92,8 @@ def save_images2(image_dir, visuals, image_path, aspect_ratio=1.0, width=256):
 
 
 SAFEHASH = [x for x in "0123456789-abcdefghijklmnopqrstuvwxyz_ABCDEFGHIJKLMNOPQRSTUVWXYZ"]
+
+
 def compress_UUID():
     '''
     根据http://www.ietf.org/rfc/rfc1738.txt,由uuid编码扩bai大字符域生成du串
@@ -108,13 +111,29 @@ def compress_UUID():
     return safe_code
 
 
-def run(
-    image,
-    model,
-    opt,
-) -> tuple[PIL.Image.Image]:
 
-    dataroot = 'images/'+compress_UUID()
+def get_68lm(imgfile, savepath, detector, predictor):
+    image = cv2.imread(imgfile)
+    rgbImg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    rects = detector(rgbImg, 1)
+    for (i, rect) in enumerate(rects):
+        landmarks = predictor(rgbImg, rect)
+        landmarks = shape_to_np(landmarks)
+        f = open(savepath, 'w')
+        for i in range(len(landmarks)):
+            lm = landmarks[i]
+            print(lm[0], lm[1], file=f)
+        f.close()
+
+
+def run(
+        image,
+        model,
+        opt,
+        detector,
+        predictor,
+) -> tuple[PIL.Image.Image,PIL.Image.Image,PIL.Image.Image,PIL.Image.Image]:
+    dataroot = 'images/' + compress_UUID()
     opt.dataroot = os.path.join(dataroot, 'src/')
     os.makedirs(opt.dataroot, exist_ok=True)
     opt.results_dir = os.path.join(dataroot, 'results/')
@@ -127,25 +146,40 @@ def run(
 
     shutil.copy(image.name, opt.dataroot)
 
-    data_loader = CreateDataLoader(opt)
-    dataset = data_loader.load_data()
+    fullname = os.path.basename(image.name)
+    name = fullname.split(".")[0]
 
-    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)
+    imgfile = os.path.join(opt.dataroot, fullname)
+    lmfile = os.path.join(opt.lm_dir, name+'.txt')
+    # 预处理数据
+    get_68lm(imgfile, lmfile, detector, predictor)
 
-    print(imgs)
-    return PIL.Image.open(imgs[0])
+    imgs = []
+    for part in ['eyel', 'eyer', 'nose', 'mouth']:
+        savepath = os.path.join(opt.bg_dir + part, name+'.png')
+        get_partmask.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)
+    #
+    # print(imgs)
+
+    return PIL.Image.open(imgs[0]),PIL.Image.open(imgs[1]),PIL.Image.open(imgs[2]),PIL.Image.open(imgs[3])
 
 
 def main():
@@ -178,14 +212,18 @@ def main():
 
     opt.checkpoints_dir = checkpoint_dir
 
-
     model = create_model(opt)
     model.setup(opt)
 
-    func = functools.partial(run, model=model, opt=opt)
+    '''
+    预处理数据
+    '''
+    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.update_wrapper(func, run)
 
-    
     gr.Interface(
         func,
         [
@@ -196,7 +234,7 @@ def main():
                 type='pil',
                 label='Result'),
         ],
-        #examples=examples,
+        # examples=examples,
         theme=args.theme,
         title=TITLE,
         description=DESCRIPTION,

requirements.txt

@@ -1,8 +1,10 @@
-torch>=0.4.0
-torchvision>=0.2.1
-dominate>=2.3.1
-visdom>=0.1.8.3
-scipy>=1.1.0
-numpy>=1.14.1
-Pillow>=5.0.0
-opencv-python>=3.4.2
+torch==1.1.0
+torchvision==0.4.0
+dominate==2.4.0
+visdom==0.1.8.9
+scipy==1.1.0
+numpy==1.16.4
+Pillow==4.3.0
+opencv-python==4.1.0.25
+dlib==19.18.0
+shapely==1.7.0