From 4dc201d6472bef3f397ce780e110d18ef026054c Mon Sep 17 00:00:00 2001 From: andreasjansson Date: Sat, 23 Oct 2021 02:02:37 +0200 Subject: [PATCH] Update Cog predictor to use a set of default target images --- predict.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/predict.py b/predict.py index 8f930dda..b93ec729 100644 --- a/predict.py +++ b/predict.py @@ -1,3 +1,5 @@ +import os +from glob import glob import cog import tempfile from pathlib import Path @@ -16,6 +18,12 @@ from insightface_func.face_detect_crop_single import Face_detect_crop as Face_detect_crop_single +TARGET_OPTIONS_DIR = "cog-targets" + +def list_target_options(): + return [os.path.basename(p).split(".")[0] for p in glob(f"{TARGET_OPTIONS_DIR}/*.jpg")] + + class Predictor(cog.Predictor): def setup(self): self.transformer_Arcface = transforms.Compose([ @@ -24,10 +32,11 @@ def setup(self): ]) @cog.input("source", type=Path, help="source image") - @cog.input("target", type=Path, help="target image") + @cog.input("target", type=str, options=list_target_options(), help="target image") @cog.input("mode", type=str, options=['single', 'all'], default='all', help="swap a single face (the one with highest confidence by face detection) or all faces in the target image") def predict(self, source, target, mode='all'): + target_path = f"{TARGET_OPTIONS_DIR}/{target}.jpg" app = Face_detect_crop_multi(name='antelope', root='./insightface_func/models') @@ -39,7 +48,7 @@ def predict(self, source, target, mode='all'): options = TestOptions() options.initialize() opt = options.parser.parse_args(["--Arc_path", 'arcface_model/arcface_checkpoint.tar', "--pic_a_path", str(source), - "--pic_b_path", str(target), "--isTrain", False, "--no_simswaplogo"]) + "--pic_b_path", target_path, "--isTrain", False, "--no_simswaplogo"]) str_ids = opt.gpu_ids.split(',') opt.gpu_ids = []