import os import time import argparse import wandb import random import tqdm import scipy as sp import cmapy as cm import scipy.interpolate as spi import numpy as np import tensorflow as tf from custom_utils.custom_utils import UV2MeshLabel, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex from plotly.subplots import make_subplots from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block from stylegan2_ae.generator import Generator from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL os.environ["CUDA_VISIBLE_DEVICES"] = "1" uv2mesh = UV2MeshLabel(res=256) def downscale_1(img): results = [img] for i in range(5): results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5)) return results def downscale_2(img): results = [img] for i in range(5): # print(np.array(img.shape[:2]) / (2 ** (i + 1))) results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True)) return results def upscale_1(imgs): results = [] for img in imgs: results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)) return results def fill(img, mask): coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0)) img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0)) img_data = img[img_idx[:,0], img_idx[:,1], :] interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) img_2 = interpolator(coords) res = np.zeros((512,256,3), dtype=np.float32) res[coords[:,0], coords[:,1], :] = img_2 res[img_idx[:,0], img_idx[:,1], :] = img_data return res def fill_uv_maps(folder, type): os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True) mask = np.load('custom_utils/mask_256x512.npy') frames = os.listdir(os.path.join(folder, type)) for frame in tqdm.tqdm(frames): id, _ = frame.split('_') # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask'] img = np.load(os.path.join(folder, type, frame)); img = img['uv_body'] filled = fill(img, mask) np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_body_posed = filled) if __name__ == '__main__': # fill_uv_maps(folder='/home/workspace/data/dataset/stylegan_jumpsuit/train', type='uv_body_posed') # pass # exit(0) allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) smpl = SMPL() factory = DatasetFactory() signature_dynamic = { 'id': (), 'frame_nr': (), 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), # 'uv_body_posed': (3, 512, 256,), # 'mesh_body_posed': (14475, 3,), 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), # 'uv_unposed': (3, 512, 256,) 'uv_curvature': (512, 256, 7), 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), # 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=1, signature=signature_dynamic) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static) step = 0 for batch in dataset_train: # batch = batch_to_dict(batch, signature_dynamic) batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) img = tf.squeeze(batch['uv_curvature']).numpy() coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0)) img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0)) img_data = img[img_idx[:,0], img_idx[:,1]] interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) img_2 = interpolator(coords) res = np.zeros((512,256)) res[coords[:,0], coords[:,1]] = img_2 res[img_idx[:,0], img_idx[:,1]] = img_data img_2 = res i = 0 down_curv = tf.image.resize(img_2[:,:,tf.newaxis], size=np.array(img_2.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) down_curv = tf.image.resize(down_curv, size=(512, 256), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)[:,:,0] img_posed = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy() coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0)) img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0)) img_data = img_posed[img_idx[:,0], img_idx[:,1], :] interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) img_posed_2 = interpolator(coords) res = np.zeros((512,256,3)) res[coords[:,0], coords[:,1], :] = img_posed_2 res[img_idx[:,0], img_idx[:,1], :] = img_data img_posed_2 = res mesh_posed = uv2mesh(tf.transpose(img_posed_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] mask = batch['mesh_mask'].numpy().astype(bool)[0] mesh_1 = uv2mesh(img_2[tf.newaxis,:,:])[0].numpy()[mask] mesh_2 = uv2mesh(down_curv[tf.newaxis,:,:])[0].numpy()[mask] F = np.load("custom_utils/smpl_faces.npy") _, faces = _maskFaces(np.ones((14475, 3)), F, mask) fig = make_subplots(rows=1, cols=2, start_cell="top-left") fig.add_trace(ex.imshow((img_2) / 6).data[0], row=1, col=1) fig.add_trace(ex.imshow((down_curv ) / 6).data[0], row=1, col=2) fig.show() mesh = tf.squeeze(batch['mesh_posed']).numpy()[mask] colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) v_color_1 = np.zeros((mesh.shape[0], 3)) v_color_2 = np.zeros((mesh.shape[0], 3)) for i in range(6+1): v_color_1[mesh_1 == i, :] = colors[i] v_color_2[mesh_2 == i, :] = colors[i] mesh_copy = mesh.copy() mesh_copy[:,1] += 1 fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=mesh[:,0], y=mesh[:,1], z=mesh[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2], vertexcolor=v_color_1, text='original' ), go.Mesh3d( # 8 vertices of a cube x=mesh_copy[:,0], y=mesh_copy[:,1], z=mesh_copy[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2], vertexcolor=v_color_2, text='rescaled' ) ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,3],), zaxis = dict(nticks=4, range=[-1,1],),) ) fig.show() print()