from PIL import Image import os import os.path as osp import numpy as np from tqdm import tqdm from einops import reduce import click import cv2 import sys sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__)))) from utils.io_utils import * from live2d.scrap_model import Live2DScrapModel, compose_from_drawables, load_detected_character, init_drawable_visible_map, Drawable from utils.visualize import pil_draw_text, visualize_segs, VALID_FACE_GROUPS, FACE_LABEL2NAME, visualize_facedet_output, LEFT_EYEBROW, RIGHT_EYEBROW, show_factorization_on_image from utils.cv import mask2rle, rle2mask @click.group() def cli(): """live2d scripts. """ @cli.command('smug') @click.option('--exec_list') @click.option('--rank_to_worldsize', default=None) def smug(exec_list, rank_to_worldsize): exec_list = load_exec_list(exec_list, rank_to_worldsize=rank_to_worldsize) import smug from smug.utils.logging import set_logger from smug import KeyPointsController, IRIAMSeparationController controller_kp = KeyPointsController() controller_sep = IRIAMSeparationController() for p in tqdm(exec_list): lmodel = Live2DScrapModel(p, pad_to_square=False, crop_to_final=False) if osp.isfile(p): src_dir = osp.dirname(p) else: src_dir = p src_img = lmodel.final[..., :3] faceparsing_output = osp.join(src_dir, 'face_parsing') os.makedirs(faceparsing_output, exist_ok=True) try: result = controller_kp.run(src_img) landmark_dict = { 'pose': result.pose, 'face_landmarks': {'left': result.face_landmarks.left, 'top': result.face_landmarks.top, 'points': result.face_landmarks.points} } dict2json(landmark_dict, osp.join(faceparsing_output, 'pose.json')) ls = controller_sep.run(src_img, {'output_folder': faceparsing_output}) except Exception as e: print(f'failed to process {p}: {e}') continue @cli.command('further_extr') @click.option('--exec_list') @click.option('--rank_to_worldsize', default=None) def further_extr(exec_list, rank_to_worldsize): exec_list = load_exec_list(exec_list, rank_to_worldsize=rank_to_worldsize) eye_mesh_dict = { '1-2-3-2-2+eyebgs-l': 'eyebgsl', '1-2-3-3-2+irides-l': 'iridesl', '1-2-3-1-2+eyelashs-l': 'eyelashsl', '1-2-3-2-1+eyebgs-r': 'eyebgsr', '1-2-3-3-1+irides-r': 'iridesr', '1-2-3-1-1+eyelashs-r': 'eyelashsr', '1-2-1-1+eyebrows-r': 'eyebrowr', '1-2-1-2+eyebrows-l': 'eyebrowl' } for p in tqdm(exec_list): try: fp = osp.join(p, 'face_parsing') parts_dict_exist = osp.exists(osp.join(fp, 'parts.json')) if parts_dict_exist: parts = json2dict(osp.join(fp, 'parts.json')) eye_parts = {} for k, n in eye_mesh_dict.items(): imgp = osp.join(fp, k + '.png') if not osp.exists(imgp) or not parts_dict_exist: eye_parts[n] = {'area': 0} continue img = np.array(Image.open(osp.join(fp, k + '.png'))) pd = parts[k] x, y, w, h = pd['x'], pd['y'], pd['w'], pd['h'] mask = img[..., -1] > 15 rect = cv2.boundingRect(cv2.findNonZero(mask.astype(np.uint8))) xyxy = [x, y, x + w, y + h] rect = [rect[0] + x, rect[1] + y, rect[2], rect[3]] rect[2] += rect[0] rect[3] += rect[1] eye_parts[n] = { 'img': img, 'xyxy': xyxy, 'mask': mask, 'rect': rect, 'area': np.sum(mask) } lmodel = Live2DScrapModel(p, pad_to_square=False, crop_to_final=False) lmodel.init_drawable_visible_map() lmodel.load_body_parsing() max_d = len(lmodel.drawables) + 1 face_max_idx = -1 face_min_idx = max_d neck_max_idx = -1 neck_min_idx = max_d nose_max_idx = -1 nose_min_idx = max_d mouth_max_idx = -1 mouth_min_idx = max_d for d in lmodel.drawables: if d.body_part_tag == 'nose': nose_max_idx = max(d.idx, nose_max_idx) nose_min_idx = min(d.idx, nose_min_idx) elif d.body_part_tag == 'mouth': mouth_max_idx = max(d.idx, mouth_max_idx) mouth_min_idx = min(d.idx, mouth_min_idx) for d in lmodel.drawables: if d.body_part_tag == 'face': tgt_max_idx = min(d.idx, nose_min_idx, mouth_min_idx) face_max_idx = max(tgt_max_idx, face_max_idx) face_min_idx = min(d.idx, face_min_idx) hair_split_idx = face_max_idx for d in lmodel.drawables: if d.body_part_tag is None or 'hair' not in d.body_part_tag: continue if d.idx > hair_split_idx: d.body_part_tag = 'hairf' else: d.body_part_tag = 'hairb' eyel_xyxy = [lmodel.final.shape[1], lmodel.final.shape[0], 0, 0] eyer_xyxy = [lmodel.final.shape[1], lmodel.final.shape[0], 0, 0] for k in {'iridesl', 'eyebgsl', 'eyelashsl'}: if 'xyxy' in eye_parts[k]: eyel_xyxy[0] = min(eyel_xyxy[0], eye_parts[k]['xyxy'][0]) eyel_xyxy[1] = min(eyel_xyxy[1], eye_parts[k]['xyxy'][1]) eyel_xyxy[2] = max(eyel_xyxy[2], eye_parts[k]['xyxy'][2]) eyel_xyxy[3] = max(eyel_xyxy[3], eye_parts[k]['xyxy'][3]) for k in {'iridesr', 'eyebgsr', 'eyelashsr'}: if 'xyxy' in eye_parts[k]: eyer_xyxy[0] = min(eyer_xyxy[0], eye_parts[k]['xyxy'][0]) eyer_xyxy[1] = min(eyer_xyxy[1], eye_parts[k]['xyxy'][1]) eyer_xyxy[2] = max(eyer_xyxy[2], eye_parts[k]['xyxy'][2]) eyer_xyxy[3] = max(eyer_xyxy[3], eye_parts[k]['xyxy'][3]) for d in lmodel.drawables: if d.body_part_tag != 'eyes': continue eye_tag = None score = 0. eye_scores = {} for ek, ed in eye_parts.items(): if ed['area'] == 0: eye_scores[ek] = [None] * 4 continue mask = ed['mask'] area, u_area, i_area = d.mask_union_intersection(mask, ed['xyxy'], final_vis_mask=True) eye_scores[ek] = [area, u_area, i_area, ed['area']] irides_scores, bg_scores = None, None eyelash_scores = eyebrow_scores = None if eye_scores['iridesl'][0] is not None: irides_scores = eye_scores['iridesl'] bg_scores = eye_scores['eyebgsl'] elif eye_scores['iridesr'][0] is not None: irides_scores = eye_scores['iridesr'] bg_scores = eye_scores['eyebgsr'] if eye_scores['eyelashsr'][0] is not None: eyelash_scores = eye_scores['eyelashsr'] elif eye_scores['eyelashsl'][0] is not None: eyelash_scores = eye_scores['eyelashsl'] if eye_scores['eyebrowr'][0] is not None: eyebrow_scores = eye_scores['eyebrowr'] elif eye_scores['eyebrowl'][0] is not None: eyebrow_scores = eye_scores['eyebrowl'] iou_i = iou_b = iou_l = iou_br = -1 scores = {'irides': 0, 'eyebg': 0, 'eyelash': 0, 'eyebrow': 0} if irides_scores is not None and bg_scores is not None and irides_scores[2] > 0 and bg_scores[2] > 0: scores['irides'] = irides_scores[2] / irides_scores[1] scores['eyebg'] = bg_scores[2] / bg_scores[1] if eyelash_scores is not None and eyelash_scores[2] > 0: scores['eyelash'] = eyelash_scores[2] / eyelash_scores[1] if eyebrow_scores is not None and eyebrow_scores[2] > 0: scores['eyebrow'] = eyebrow_scores[2] / eyebrow_scores[1] k = max(scores, key=scores.get) def rect_include(xyxy1, dict2): if 'xyxy' not in dict2: return False xyxy2 = dict2['xyxy'] return xyxy1[0] > xyxy2[0] and xyxy1[1] > xyxy2[1] and xyxy1[2] < xyxy2[2] and xyxy1[3] < xyxy2[3] if scores[k] > 0: d.body_part_tag = k else: x1, y1, x2, y2 = d.xyxy y = (y1 + y2) / 2 if y < eyel_xyxy[1] or y < eyer_xyxy[1]: d.body_part_tag = 'eyebrow' elif rect_include(d.xyxy, eye_scores['iridesl']) or rect_include(d.xyxy, eye_scores['iridesr']): d.body_part_tag = 'irides' elif rect_include(d.xyxy, eye_scores['eyebgsl']) or rect_include(d.xyxy, eye_scores['eyebgsr']): d.body_part_tag = 'eyebg' else: d.body_part_tag = 'eyelash' lmodel.save_body_parsing(save_name='bodyparsingv3') # hairf = lmodel.compose_bodypart_drawables('hairf') # hairb = lmodel.compose_bodypart_drawables('hairb') # irides = lmodel.compose_bodypart_drawables('irides') # eyebg = lmodel.compose_bodypart_drawables('eyebg') # eyelash = lmodel.compose_bodypart_drawables('eyelash') # eyebrow = lmodel.compose_bodypart_drawables('eyebrow') # save_tmp_img( # imglist2imgrid([lmodel.final, hairf, hairb, irides, eyebg, eyelash, eyebrow], fix_size=512) # ) # pass except Exception as e: # raise print(f'failed to process {p}: {e}') continue if __name__ == '__main__': cli()