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