inference/scripts/vis_texture.py

import sys
import os.path as osp
import os
import argparse
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))

default_n_threads = 8
os.environ['OPENBLAS_NUM_THREADS'] = f"{default_n_threads}"
os.environ['MKL_NUM_THREADS'] = f"{default_n_threads}"
os.environ['OMP_NUM_THREADS'] = f"{default_n_threads}"

from tqdm import tqdm
import numpy as np
import cv2
from PIL import Image
from utils.io_utils import *
from talking_head.mesh import *
from utils.torchcv import cluster_inpaint_part
from utils.visualize import pil_draw_text
from utils.cv import *
# from utils.torch_utils import seed_everything

def further_extr(srcd: str, rotate=True):


    def _label_lr_split(labels, stats, id1, id2):
        label1 = (labels == id1).astype(np.uint8) * 255
        label2 = (labels == id2).astype(np.uint8) * 255

        stats1, stats2 = stats[id1], stats[id2]

        x1 = stats[id1][0] + stats[id1][2] / 2
        x2 = stats[id2][0] + stats[id2][2] / 2

        if x2 < x1:
            return label2, label1, stats2, stats1
        else:
            return label1, label2, stats1, stats2


    def _save_part(part_name):
        if part_name not in tag2pinfo:
            print(f'{part_name} is not valid')
            return
        part_dict = tag2pinfo[part_name]
        img = part_dict.pop('img')
        if 'mask' in part_dict:
            part_dict.pop('mask')

        mask = img[..., -1] > 10

        xywh = cv2.boundingRect(cv2.findNonZero(mask.astype(np.uint8)))
        xyxy = np.array(xywh).copy()
        xyxy[2] += xyxy[0]
        xyxy[3] += xyxy[1]

        depth = part_dict.pop('depth')
        depth_median = np.median(depth[mask])
        depth = depth[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]]
        img = img[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]]

        # dmin, dmax = np.min(depth), np.max(depth)
        # depth = (depth - dmin) / (dmax - dmin + 1e-8) * 255
        depth = np.clip(depth, 0, 1) * 255
        depth = np.round(depth).astype(np.uint8)

        x1, y1, x2, y2 = part_dict['xyxy']
        part_dict['xyxy'] = [x1 + xyxy[0], y1 + xyxy[1], x1 + xyxy[2], y1 + xyxy[3]]
        # part_dict['depth_min'] = dmin
        # part_dict['depth_max'] = dmax
        part_dict['depth_median'] = depth_median
        part_dict['tag'] = part_name

        Image.fromarray(img).save(osp.join(saved, part_name + '.png'))
        Image.fromarray(depth).save(osp.join(saved, part_name + '_depth.png'))


    def _process_cuts(img, depth, src_xyxy, tgt_bbox, p=5):
        tx1, ty1, tx2, ty2 = tgt_bbox[:4]
        tx2 += tx1
        ty2 += ty1
        img = img[ty1: ty2, tx1: tx2]
        depth = depth[ty1: ty2, tx1: tx2]
        fxyxy = [tx1 + src_xyxy[0], ty1 + src_xyxy[1], tx2 + src_xyxy[0], ty2 + src_xyxy[1]]

        return img, depth, fxyxy


    saved = osp.join(srcd, 'optimized')
    infos = json2dict(osp.join(srcd, 'info.json'))
    os.makedirs(saved, exist_ok=True)

    fullpage, infos, part_dict_list = load_parts(srcd, rotate=rotate)

    # optim_depth(part_dict_list, fullpage)

    tag2pinfo = {}
    for pinfo in part_dict_list:
        tag = pinfo['tag']
        tag2pinfo[tag] = pinfo

    # part_info = tag2pinfo.pop('eyes')
    # num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
    #     part_info['mask'].astype(np.uint8) * 255, connectivity=8)
    # if len(stats) > 2:
    #     stats = np.array(stats)
        
    #     if len(stats[..., -1]) >= 5:
    #         stats_order = np.argsort(stats[..., -1])[::-1][1:]
    #         eyel_mask, eyer_mask, statsl, statsr = _label_lr_split(labels, stats, stats_order[0], stats_order[1])
    #         img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsl)
    #         tag2pinfo['eyer'] = {'img': img, 'xyxy': xyxy, 'depth': depth}
    #         img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsr)
    #         tag2pinfo['eyel'] = {'img': img, 'xyxy': xyxy, 'depth': depth}

    #         browl_mask, browr_mask, statsl, statsr = _label_lr_split(labels, stats, stats_order[2], stats_order[3])
    #         img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsl)
    #         tag2pinfo['browr'] = {'img': img, 'xyxy': xyxy, 'depth': depth}
    #         img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsr)
    #         tag2pinfo['browl'] = {'img': img, 'xyxy': xyxy, 'depth': depth}

    if 'handwear' in tag2pinfo:
        part_info = tag2pinfo.pop('handwear')
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            part_info['mask'].astype(np.uint8) * 255, connectivity=8)
        if len(stats) > 2:
            stats = np.array(stats)
            stats_order = np.argsort(stats[..., -1])[::-1][1:]
            arml_mask, armr_mask, statsl, statsr = _label_lr_split(labels, stats, stats_order[0], stats_order[1])
            img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsl)
            tag2pinfo['arml'] = {'img': img, 'xyxy': xyxy, 'depth': depth, 'mask': img[..., -1] > 10}
            img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsr)
            tag2pinfo['armr'] = {'img': img, 'xyxy': xyxy, 'depth': depth}
        else:
            tag2pinfo['handwear'] = part_info

    if 'ears' in tag2pinfo:
        part_info = tag2pinfo.pop('ears')
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            part_info['mask'].astype(np.uint8) * 255, connectivity=8)
        if len(stats) > 2:
            stats = np.array(stats)
            stats_order = np.argsort(stats[..., -1])[::-1][1:]
            arml_mask, armr_mask, statsl, statsr = _label_lr_split(labels, stats, stats_order[0], stats_order[1])
            img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsl)
            tag2pinfo['earr'] = {'img': img, 'xyxy': xyxy, 'depth': depth}
            img, depth, xyxy = _process_cuts(part_info['img'], part_info['depth'], part_info['xyxy'], statsr)
            tag2pinfo['earl'] = {'img': img, 'xyxy': xyxy, 'depth': depth}
        else:
            tag2pinfo['ears'] = part_info

    # if 'headwear' in tag2pinfo:
    #     part_info = tag2pinfo.pop('headwear')
    #     tag2pinfo['hair']['img'] = img_alpha_blending([tag2pinfo['hair'], part_info], xyxy=tag2pinfo['hair']['xyxy'], premultiplied=False)
    # if 'headwear' in tag2pinfo:
    #     part_info = tag2pinfo.pop('headwear')
    #     tag2pinfo['hair']['img'] = img_alpha_blending([tag2pinfo['hair'], part_info], xyxy=tag2pinfo['hair']['xyxy'], premultiplied=False)

    # if 'footwear' in tag2pinfo:
    #     part_info = tag2pinfo.pop('footwear')
    #     tag2pinfo['legwear']['img'] = img_alpha_blending([tag2pinfo['legwear'], part_info], xyxy=tag2pinfo['legwear']['xyxy'], premultiplied=False)

    if 'hair' in tag2pinfo:
        part_info = tag2pinfo.pop('hair')
        parts = cluster_inpaint_part(**part_info)
        parts.sort(key=lambda x: x['depth_median'])
        tag2pinfo['hairf'] = parts[0]
        tag2pinfo['hairb'] = parts[1]

    if 'bottomwear' in tag2pinfo:
        part_info = tag2pinfo.pop('bottomwear')
        parts = cluster_inpaint_part(**part_info)
        parts.sort(key=lambda x: x['depth_median'])
        tag2pinfo['bottomwearf'] = parts[0]
        tag2pinfo['bottomwearb'] = parts[1]

    if 'topwear' in tag2pinfo:
        part_info = tag2pinfo.pop('topwear')
        parts = cluster_inpaint_part(**part_info)
        parts.sort(key=lambda x: x['depth_median'])
        tag2pinfo['topwearf'] = parts[0]
        tag2pinfo['topwearb'] = parts[1]

    if 'handwear' in tag2pinfo:
        part_info = tag2pinfo.pop('handwear')
        parts = cluster_inpaint_part(**part_info)
        parts.sort(key=lambda x: x['depth_median'])
        tag2pinfo['handwearf'] = parts[0]
        tag2pinfo['handwearb'] = parts[1]

    # if 'footwear' in tag2pinfo:
    #     tag2pinfo.pop('footwear')

    if 'nose' in tag2pinfo:
        xyxy = tag2pinfo['nose']['xyxy']
        tag2pinfo['nose']['img'][..., :3] = fullpage[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2], :3]

    if 'mouth' in tag2pinfo:
        xyxy = tag2pinfo['mouth']['xyxy']
        tag2pinfo['mouth']['img'][..., :3] = fullpage[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2], :3]

    for t in tag2pinfo:
        _save_part(t)

    if 'face' in tag2pinfo:
        for t in ['nose', 'mouth', 'eyes']:
            if t in tag2pinfo:
                if tag2pinfo[t]['depth_median'] > tag2pinfo['face']['depth_median']:
                    tag2pinfo[t]['depth_median'] = tag2pinfo['face']['depth_median'] - 0.001
        for t in ['earr', 'earl', 'ears']:
            if t in tag2pinfo:
                tag2pinfo[t]['depth_median'] = tag2pinfo['face']['depth_median'] + 0.001

    frame_size = fullpage.shape[:2]

    dict2json({'parts': tag2pinfo, 'frame_size': frame_size}, osp.join(saved, 'info.json'))
    pass



def texture_partition(mesh_list: list, composed):
    left_out = composed['left_out']
    left_out_updated = None
    if left_out is not None and left_out[2] > 0 and left_out[3] > 0:
        lf_space = left_out[2] * left_out[3]
        x, y, w, h = left_out
        upd_idx = None
        expire_counter = 0
        for ii, m in enumerate(mesh_list):
            ih, iw = m['img'].shape[:2]
            if iw <= w and ih <= h:
                if (w-iw) * h > (h - ih) * w:
                    lf_updated = [iw + x, y, w-iw, h]
                else:
                    lf_updated = [x, y + ih, w, h - ih]
                sp_updated = lf_updated[2] * lf_updated[3]
                if sp_updated < lf_space:
                    lf_space = sp_updated
                    upd_idx = ii
                    left_out_updated = lf_updated
            else:
                expire_counter += 1
                if expire_counter >= 4:
                    break
                
                
    nw_counter = composed['nw_counter']
    if left_out_updated is not None:
        mesh = mesh_list.pop(upd_idx)
        mh, mw = mesh['img'].shape[:2]
        # composed['bx'] += mw
        # composed['by'] += mh
        mesh['compose_pos'] = composed['left_out'][:2]
        composed['left_out'] = left_out_updated
        composed['mesh_list'].append(mesh)
    else:
        bx , by = composed['bx'], composed['by']
        mesh = mesh_list.pop(0)
        mh, mw = mesh['img'].shape[:2]
        if (bx + mw) * by < (by + mh) * bx or nw_counter > 0:
            composed['bx'] += mw
            mesh['compose_pos'] = [bx, 0]
            if by - mh > 0:
                left_out_updated = [bx, mh, mw, by - mh]
            else:
                left_out_updated = [0, by, bx, mh - by]
                composed['by'] = mh
            
        else:
            nw_counter += 1
            composed['by'] += mh
            mesh['compose_pos'] = [0, by]
            if bx - mw > 0:
                left_out_updated = [mw, by, bx - mw, mh]
            else:
                left_out_updated = [bx, 0, mw - bx, by]
                composed['bx'] = mw
            
        composed['mesh_list'].append(mesh)
        composed['left_out'] = left_out_updated
    
    composed['nw_counter'] = nw_counter
        
    if len(mesh_list) > 0:
        texture_partition(mesh_list, composed)




def texture_mosaic(srcd):
            
    src_infop = osp.join(srcd, 'info.json')
    # src_info = json2dict(src_infop)

    rotate = False
    fullpage, src_info, part_dict_list = load_parts(srcd, pad=[32, 0, 0, 0], rotate=rotate)
    ih, iw = fullpage.shape[:2]
    fp = None
    lp = None
    for p in part_dict_list:
        if p['tag'] == 'footwear':
            fp = p
        elif p['tag'] == 'legwear':
            lp = p


    if fp is not None and lp is not None:
        ksize = 4
        element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * ksize + 1, 2 * ksize + 1),(ksize, ksize))
        

        img1 = np.full((ih, iw, 4), 0, dtype=np.uint8)
        img2 = img1.copy()
        xyxy = lp['xyxy']
        d = np.full((ih, iw), 0, dtype=np.float32)
        img1[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]] = lp['img']
        d[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]] = lp['depth']
        xyxy = fp['xyxy']
        img2[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]] = fp['img']
        mask = img2[..., -1]
        mask = cv2.dilate(mask, element)
        img1[..., -1] = np.bitwise_and(img1[..., -1], np.bitwise_not(mask))
        d = 1-(1-d) * (img1[..., -1]) / 255.
        xyxy = cv2.boundingRect(cv2.findNonZero((img1[..., -1] > 127).astype(np.uint8)))
        xyxy = list(xyxy)
        xyxy[2] += xyxy[0]
        xyxy[3] += xyxy[1]
        lp['img'] = img1[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]]
        lp['depth'] = d[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]]
        lp['xyxy'] = xyxy
        


    mesh_list = part_dict_list
    mesh_list.sort(key=lambda x: x['img'].shape[0], reverse=True)
    composed= {'left_out': None, 'mesh_list': [], 'bx': 0, 'by': 0, 'nw_counter': 0}
    texture_partition(mesh_list, composed)
    texture = np.zeros((composed['by'], composed['bx'], 4), dtype=np.uint8)
    depth_texture = np.full((composed['by'], composed['bx']), 1, dtype=np.float32)
    UVs = []
    pos = []
    # composed['mesh_list'].sort(key=lambda x: x['depth_median'], reverse=True)

    for ii, mesh in enumerate(composed['mesh_list']):
        img = mesh['img']
 
        depth = mesh['depth']
        x, y = mesh['compose_pos']
        h, w = img.shape[:2]
        texture[y: y + h, x: x + w] = img
        depth_texture[y: y + h, x: x + w] = depth

    texture = np.ascontiguousarray(texture)
    h, w = texture.shape[:2]
    if rotate:
        texture = np.rot90(texture, 1)
    texture_walpha = texture.copy()
    # texture = checkerboard_vis(np.array(texture))
    texture = Image.fromarray(texture)
    tlist = []
    plist = []
    for ii, mesh in enumerate(composed['mesh_list']):
        tag = mesh['tag']
        mh, mw = mesh['img'].shape[:2]
        x, y = mesh['compose_pos']
        nx = x
        ny = y
        if rotate:
            nx = y
            ny = w - x - mw
        # nx = x
        # ny = y
        pos = (nx, ny)
        plist.append(pos)
        tlist.append(tag)
    pil_draw_text(texture, tlist, plist, font_size=32)

    import matplotlib as mpl
    import matplotlib.pyplot as plt

    colormap = plt.colormaps['inferno']

    if rotate:
        depth_texture = np.rot90(depth_texture, 1)
    depth_vis = colormap(1-depth_texture, bytes=True)

    # save_tmp_img(depth_vis, 'local_tst_depth.png')
    # save_tmp_img(texture)
    return texture_walpha, texture, depth_vis

    # return mesh_compose

srcp = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output114514/gyrojeff_girl_tachie_no_glass_full_body_mizuki_niji_7_0aa1ca73_7f25'
srcp = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output114514/gyrojeff_boy_tachie_no_glass_full_body_hair_accessories_nij_25115854/src_img.png'
srcp = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output114514/gyrojeff_girl_tachie_no_glass_full_body_mizuki_niji_7_0aa1ca73_7f25 (2)/src_img.png'
srcp = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output0/gyrojeff_girl_tachie_no_glass_full_body_mizuki_niji_7_0aa1ca73_7f25 (3)/src_img.png'
srcp = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output114514/gyrojeff_girl_tachie_no_glass_full_body_mizuki_niji_7_fd85b9a1_09b2/src_img.png'

srcd = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output114514'
srcd = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output0'
srcd = 'workspace/datasets/testcaseall_output'
srcd = 'tmp/manga_lovehina'
# srcd = 'workspace/datasets/l2d_eval_oa/l2d_eval2_output42'
save_dir = osp.join('tmp', osp.basename(srcd) + '_texture')
os.makedirs(save_dir, exist_ok=True)

for d in tqdm(os.listdir(srcd)):
    if 'lovehina_seed105' not in d:
        continue
    imgn = osp.basename(d)
    srcp = osp.join(srcd, d)
    
    op_dir = osp.join(srcp, 'optimized')

    src_img = np.array(Image.open(osp.join(srcp, 'src_img.png')))

    if not osp.exists(op_dir):
        further_extr(srcp, rotate=False)

    savename = osp.join(save_dir, d)
    texture_walpha, texture, depth_vis = texture_mosaic(srcp)
    save_tmp_img(texture, savename+'_texture.png')
    save_tmp_img(texture_walpha, savename+'_texture_wocheckerboard.png')
    save_tmp_img(depth_vis, savename+'_depth.png')

    src_infop = osp.join(op_dir, 'info.json')
    src_info = json2dict(src_infop)
    load_img_depth(op_dir, src_info, pad=0)

    img_list = []
    for t, pd in src_info['parts'].items():
        if 'xyxy' not in pd:
            pd['xyxy'] = [0, 0, pd['img'].shape[1], pd['img'].shape[0]]
        if 'depth_median' not in pd:
            pd['depth_median'] = np.median(pd['depth'][pd['mask'] > 127])
        pd['tag'] = t
        pd.pop('depth')
        img_list.append({'img': pd['img'], 'xyxy': pd['xyxy'], 'layer_name': pd['tag'], 'depth_median': pd['depth_median']})
    img_list.sort(key=lambda x: x['depth_median'], reverse=True)

    sz = src_img.shape[:2]

    xyxy = np.array(cv2.boundingRect(cv2.findNonZero(src_img[..., -1])))
    xyxy[[2, 3]] += xyxy[[0, 1]]
    src_img = src_img[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]].copy()

    img = img_alpha_blending(
        img_list, final_size=sz, premultiplied=False)

    save_tmp_img(img[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]], savename+'_recon.png')
    save_tmp_img(src_img, savename+'_input.png')
    # break