feng2022's picture
Time_TravelRephotography
47c46ea
raw
history blame
1.45 kB
from argparse import ArgumentParser
import os
from os.path import join as pjoin
from subprocess import run
import numpy as np
import cv2
from tqdm import tqdm
def create_skin_mask(anno_dir, mask_dir, skin_thresh=13, include_hair=False):
names = os.listdir(anno_dir)
names = [n for n in names if n.endswith('.png')]
os.makedirs(mask_dir, exist_ok=True)
for name in tqdm(names):
anno = cv2.imread(pjoin(anno_dir, name), 0)
mask = np.logical_and(0 < anno, anno <= skin_thresh)
if include_hair:
mask |= anno == 17
cv2.imwrite(pjoin(mask_dir, name), mask * 255)
def main(args):
FACE_PARSING_DIR = 'third_party/face_parsing'
main_env = os.getcwd()
os.chdir(FACE_PARSING_DIR)
tmp_parse_dir = pjoin(args.out_dir, 'face_parsing')
cmd = [
'python',
'test.py',
args.in_dir,
tmp_parse_dir,
]
print(' '.join(cmd))
run(cmd)
create_skin_mask(tmp_parse_dir, args.out_dir, include_hair=args.include_hair)
os.chdir(main_env)
def parse_args(args=None, namespace=None):
parser = ArgumentParser("Face Parsing and generate skin (& hair) mask")
parser.add_argument('in_dir')
parser.add_argument('out_dir')
parser.add_argument('--include_hair', action="store_true", help="include hair in the mask")
return parser.parse_args(args=args, namespace=namespace)
if __name__ == "__main__":
main(parse_args())