inference/scripts/further_extr.py

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()