From 86136338a1036074bcdc75d1426f1be1435ded33 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Sep 2022 21:56:08 +0000 Subject: [PATCH] original unpose sript --- unpose.py | 388 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 388 insertions(+) create mode 100644 unpose.py diff --git a/unpose.py b/unpose.py new file mode 100644 index 0000000..f5f92d6 --- /dev/null +++ b/unpose.py @@ -0,0 +1,388 @@ +import os +import sys +import pickle +import time +import json +import math +import wandb +import argparse + +import numpy as np +import tensorflow as tf +import scipy.io as sio +import scipy as sp +import matplotlib.pyplot as plt +import plotly.graph_objects as go +from scipy.spatial.transform import Rotation + +from tqdm import tqdm +from datetime import datetime + + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from FeatureLines.superSMPL.IO import readPC2, writePC2 +from Data.Dataloader import DataLoader +from Models.Reconstructor import Reconstructor_v2 as Reconstructor +from Models.FeatureLineGenerator import FeatureLineGenerator +from Models.losses import L2_loss, bend_loss, collision_loss, normal_loss, L1_loss, feature_lines_loss, edge_loss, normal_loss_v2 +# from Models.metrics import L2_error, surf2surf, feature_lines_error +from Models.utils import maskFacesBatch, pretty_json, repeat_faces, repeat_laplacian, sparse_to_tensor, getFlattenedSortedUniqueFeatureLineIndices, _maskFaces +from FeatureLines.superSMPL.smpl import SMPL +from FeatureLines.superSMPL.smpl_2 import SMPL as SMPL_2 + + + +os.environ["CUDA_VISIBLE_DEVICES"] = "6" +gpus = tf.config.experimental.list_physical_devices('GPU') +if gpus: + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + + except RuntimeError as e: + print(e) + + +# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='unpose', name='new_smpl_2', notes='') +# artifact = wandb.Artifact('code', type='code') +# artifact.add_file('Data/unpose.py') +# artifact.add_file('Data/verify_unposed.py') +# artifact.add_file('FeatureLines/superSMPL/smpl.py') +# artifact.add_file('FeatureLines/superSMPL/smpl_2.py') +# run.log_artifact(artifact) + + +def read_garment(folder, id): + garment_dir = os.path.join(folder, id) + + garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')] + assert len (garments) == 1 + + garment = garments[0] + + posed = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + pose, shape, gender = info['poses'].astype(np.float32), info['shape'], info['gender'] + shape = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) + mask = np.load(os.path.join(folder, id, garment + '_mask.npy')) + offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + template = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B'] + mesh_static = (template + offsets) * mask[:, np.newaxis] + + posed, mesh_static = posed[:, mask, :], mesh_static[mask, :] + mask_idx = np.array(range(14475))[mask] + + return posed, mask, pose, shape, mesh_static, mask_idx + +def plot_3d(mesh_orig, mesh_pred, unposed, faces): + _mesh_pred = mesh_pred.copy() + _unposed = unposed.copy() + _mesh_pred[:,1] += 1 + _unposed[:,1] += 3 + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh_orig[:,0], + y=mesh_orig[:,1], + z=mesh_orig[:,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_pred[:,0], + y=_mesh_pred[:,1], + z=_mesh_pred[:,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=_unposed[:,0], + y=_unposed[:,1], + z=_unposed[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + # vertexcolor=collision_color + ) + ]) + + 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],),) + ) + + return fig + + + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range") + parser.add_argument("-s", "--start", help = "start of range") + parser.add_argument("-e", "--end", help = "end of range") + parser.add_argument("-n", "--name", help = "end of range") + args = parser.parse_args() + + F = np.load('FeatureLines/superSMPL/smpl_faces.npy') + folder = args.folder + smpl = SMPL('FeatureLines/superSMPL') + smpl_2 = SMPL_2('FeatureLines/superSMPL') + iterations = 10000 + + ids = os.listdir(folder)[int(args.start):int(args.end)] + + for id in tqdm(ids): + if not os.path.isdir(os.path.join(folder, id)): + continue + + os.makedirs(os.path.join(folder, id, 'mesh_unposed'), exist_ok=True) + + posed, mask, pose, shape, mesh_static, mask_idx = read_garment(folder, id) + gender = int(shape[-1] > 0) + init = Rotation.from_euler('zyz', [90,90,90], degrees=True).apply(mesh_static.copy()).astype(np.float32) + + existing_frames = os.listdir(os.path.join(folder, id, 'mesh_unposed')) + if len(existing_frames) == posed.shape[0]: + continue + + R = smpl.run_rodrigues(np.transpose(pose, [1,0])) + + frames = posed.shape[0] + + # rot = Rotation.from_euler('z', [90], degrees=True) + # rest_pose = np.zeros((1,72), dtype=np.float32) + # R_rest = smpl.run_rodrigues(rest_pose) + # body = smpl.run(R_rest, shape[tf.newaxis,:])[0].numpy()[:, mask,:] + + # # verts = smpl.run(R, tf.tile(shape[tf.newaxis,:], [frames,1]))[0].numpy()[:,mask,:] + # # verts = tf.stack([rot.apply(verts[i]).astype(np.float32) for i in range(frames)]) + # verts = tf.tile(body, [frames,1,1]) + + # start = time.perf_counter() + # result = smpl.run_garment(R, tf.tile(shape[tf.newaxis,:], [frames,1]), verts, filter_idx=mask_idx) + # orig_time = time.perf_counter() - start + # print(orig_time) + + # start = time.perf_counter() + # result_2 = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx) + # new_time = time.perf_counter() - start + # print(new_time) + + # # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4 + # # x = np.sum(test, axis=-1).sum(-1).sum(-1) + # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))).numpy() + + # print('speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error)) + + # # continue + + # batch = np.random.choice(np.array(range(frames))) + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # fig = go.Figure(data=[ + # go.Mesh3d( + # # 8 vertices of a cube + # x=result[batch,:,0], + # y=result[batch,:,1], + # z=result[batch,:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2], + # text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error), + # ), + # go.Mesh3d( + # # 8 vertices of a cube + # x=result_2[batch,:,0], + # y=result_2[batch,:,1], + # z=result_2[batch,:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2], + # text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error), + # )], + # ) + + # wandb.log({id: fig}) + + # continue + + + + + + + + + lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1000.0, decay_steps=40000, decay_rate=0.1, staircase=False) + optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) + unposed = tf.Variable(tf.tile(init[tf.newaxis,:,:], [frames,1,1])) + + + @tf.function + def skinning(verts): + result = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx) + return result + + + @tf.function + def error(): + pred = skinning(unposed) + loss = tf.reduce_mean(tf.math.abs(pred - posed)) + + return loss + + start = time.perf_counter() + euclid_error = 10000 + + + while True: + optimizer.minimize(error, [unposed]) + + i = optimizer.iterations.numpy() + if i % 1000 == 0: + pred = skinning(unposed) + # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1)) + loss = tf.reduce_mean(tf.math.abs(pred - posed)) + euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed - pred) ** 2, axis=-1))) + print('id: ' + id + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start)) + start = time.perf_counter() + + if euclid_error < 0.0001: + break + + if euclid_error < 0.0005 and i > 15000: + break + + if euclid_error < 0.001 and i > 30000: + break + + if i > 40000: + break + + # b = np.random.choice(np.array(range(frames))) + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # plot_3d(posed[b], pred.numpy()[b], unposed.numpy()[b], faces).show() + + + result = np.zeros((frames, 14475, 3), dtype = np.float32) + result[:, mask,:] = unposed.numpy() + for frame in range(frames): + np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result[frame]) + + + + continue + + + + + + + + + + offset = 1 + + + for frame in range(posed.shape[0]): + + lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(10.0, decay_steps=10000, decay_rate=0.1, staircase=False) + optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) + unposed = tf.Variable(np.random.rand(offset, posed.shape[1], 3).astype(np.float32) * 2 - 1) + # unposed = tf.Variable(init) + + @tf.function + def skinning(verts): + # verts = smpl.run(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]))[0].numpy()[:,mask,:] + + + # start = time.perf_counter() + result = smpl.run_garment(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]), verts, filter_idx=mask_idx) + # print(time.perf_counter() - start) + + # start = time.perf_counter() + # result_2 = smpl_2.run_garment(R[frame:frame+offset], shape[:10], gender, verts, filter_idx=mask_idx) + # print(time.perf_counter() - start) + + # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4 + # x = np.sum(test, axis=-1).sum(-1).sum(-1) + # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))) + + # batch = 25 + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # go.Figure(data=[ + # go.Mesh3d( + # # 8 vertices of a cube + # x=result[batch,:,0], + # y=result[batch,:,1], + # z=result[batch,:,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=result_2[batch,:,0], + # y=result_2[batch,:,1], + # z=result_2[batch,:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2], + # )]).show() + + return result + + @tf.function + def error(): + pred = skinning(unposed) + loss = tf.reduce_mean(tf.math.abs(pred - posed[frame])) + + return loss + + start = time.perf_counter() + euclid_error = 10000 + + for i in range(iterations): + optimizer.minimize(error, [unposed]) + if i % 50 == 0: + pred = skinning(unposed) + # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1)) + loss = tf.reduce_mean(tf.math.abs(pred - posed[frame])) + euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed[frame] - pred) ** 2, axis=-1))) + print('id: ' + id + ' frame: ' + str(frame) + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start)) + start = time.perf_counter() + + if euclid_error < 0.001: + break + + result = np.zeros((14475, 3), dtype = np.float32) + result[mask,:] = unposed.numpy() + np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result) + + init = unposed + + # pred, unposed = pred.numpy(), unposed.numpy() + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # k = 0 + # plot_3d(posed[k], pred[k], unposed[k], faces).show() + + # print() + + + \ No newline at end of file