import os import time import argparse import wandb import random import tqdm import cv2 import scipy as sp import scipy.interpolate as spi import numpy as np import tensorflow as tf from custom_utils.custom_utils import UV2Mesh, 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 load_models import load_generator, load_discriminator from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss # 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"] = "6" uv2mesh = UV2Mesh(res=256) def downscale_1(img): results = [img] for i in range(5): # print(str(np.array(img.shape[:2]) / (2 ** (i + 1))) + ' ' + str((2 * (i + 1)))) 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)) size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32) size = (size[1], size[0]) results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST)) 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 optimize(img, gt): lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1.0, decay_steps=40000, decay_rate=0.1, staircase=False) optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) img = tf.Variable(img) @tf.function def error(): pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2) return loss for i in range(1000): optimizer.minimize(error, [img]) if i % 100 == 0: pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2) print('it: ' + str(i) + ' loss: ' + str(loss.numpy())) return img 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': (), '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,) } 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, use_filled=True) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True) 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.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[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)) res[coords[:,0], coords[:,1], :] = img_2 res[img_idx[:,0], img_idx[:,1], :] = img_data img_2 = res uv_mask = batch['uv_mask'].numpy()[0] uv_mask = downscale_2(uv_mask) # for i in range(0,3): # uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)#.astype(bool)[:, :, np.newaxis] # for i in range(3,6): # uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3), (2, 2)), iterations=1)#.astype(bool)[:, :, np.newaxis] down_1 = downscale_1(img) down_2 = downscale_1(img_2) up_1 = upscale_1(down_1) # down_2 = [down_2[i] * uv_mask[i][:,:,np.newaxis] for i in range(len(down_2))] up_2 = upscale_1(down_2) mask = batch['mesh_mask'].numpy().astype(bool)[0] mesh_2 = [uv2mesh(tf.transpose(up_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_2))] mesh_1 = [uv2mesh(tf.transpose(up_1[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_1))] F = np.load("custom_utils/smpl_faces.npy") _, faces = _maskFaces(np.ones((14475, 3)), F, mask) # mesh_downscaled = mesh_2[-1].numpy()[mask, :] # uv_optimized = optimize(down_2[-1], tf.transpose(batch['uv_posed'][0], [1,2,0])) # uv_optimized = tf.image.resize(uv_optimized, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) # mesh_optimized = uv2mesh(tf.transpose(uv_optimized, [2,0,1])[tf.newaxis,:,:,:])[0].numpy()[mask, :] # go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=mesh_downscaled[:,0], # y=mesh_downscaled[:,1], # z=mesh_downscaled[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2] # ), # go.Mesh3d( # # 8 vertices of a cube # x=mesh_optimized[:,0], # y=mesh_optimized[:,1], # z=mesh_optimized[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2], # ) # ]).show() # fig = make_subplots(rows=2, cols=6, start_cell="top-left") # for i, d1 in enumerate(down_1): # fig.add_trace(ex.imshow((d1 + 1) * 128).data[0], row=1, col=i+1) # for i, u1 in enumerate(up_1): # fig.add_trace(ex.imshow((u1 + 1) * 128).data[0], row=2, col=i+1) # fig.show() fig = make_subplots(rows=2, cols=6, start_cell="top-left") for i, d2 in enumerate(uv_mask): fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1) for i, u2 in enumerate(up_2): fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=2, col=i+1) fig.show() fig = make_subplots(rows=2, cols=6, start_cell="top-left", specs = [ [{ 'type': 'xy' } for _ in range(6)], [{ 'type': 'scene' } for _ in range(6)] ]) for i, u2 in enumerate(up_2): fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=1, col=i+1) # for i, m1 in enumerate(mesh_1): # m1 = m1.numpy()[mask, :] # fig.add_trace(go.Mesh3d( # # 8 vertices of a cube # x=m1[:,0], # y=m1[:,1], # z=m1[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2] # ), # row=1, col=i+1) for i, m2 in enumerate(mesh_2): m2 = m2.numpy()[mask, :] fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=m2[:,0], y=m2[:,1], z=m2[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), row=2, col=i+1) fig.show() print()