Spaces:
Runtime error
Runtime error
File size: 1,449 Bytes
47c46ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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())
|