From 89fca608a67f02d7a5d8da06e70cf1e326bbc8be Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:07:13 +0000 Subject: [PATCH] uv conversion acceleration --- convert_to_uv.py | 28 +- convert_to_uv_local.py | 616 ++++++++++++++++++++++++++++++++++ convert_to_uv_local_single.py | 539 +++++++++++++++++++++++++++++ 3 files changed, 1172 insertions(+), 11 deletions(-) create mode 100644 convert_to_uv_local.py create mode 100644 convert_to_uv_local_single.py diff --git a/convert_to_uv.py b/convert_to_uv.py index ae0a904..91159b8 100644 --- a/convert_to_uv.py +++ b/convert_to_uv.py @@ -448,7 +448,7 @@ def test(verts, faces, values): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt_25_2') + parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train') parser.add_argument("-s", "--start", help = "start of range", default=0) parser.add_argument("-e", "--end", help = "end of range", default=10) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') @@ -503,7 +503,7 @@ if __name__ == '__main__': if not os.path.isdir(garment_dir): continue - # if os.path.exists(os.path.join(folder, id, 'gt_uv')): + # if os.path.exists(os.path.join(folder, id, 'uv_posed_body_no_orient')): # continue garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')] @@ -526,19 +526,21 @@ if __name__ == '__main__': offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] body_static = sio.loadmat(os.path.join(garment_dir, 'info_body.mat'))['B'] mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool)) - # body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] - gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] + # gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] # curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy')) # mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] # mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])]) - mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] - mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])]) + # mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])]) # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])]) + mesh_unposed_no_shape_one_hot = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz')) for i in range(body.shape[0])] + mesh_unposed_no_shape_one_hot = np.stack([mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(body.shape[0])]) verts_static = (body_static + offsets_static) # * mask[:, np.newaxis] - verts_dynamic = gt + # verts_dynamic = gt # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing? @@ -608,7 +610,8 @@ if __name__ == '__main__': # uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) # uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) # uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) - uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) # uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) print(time.perf_counter() - start) @@ -637,10 +640,13 @@ if __name__ == '__main__': # for i in range(uv_unposed.shape[0]): # np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :]) - os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True) - for i in range(uv_unposed_no_shape.shape[0]): - np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :]) + # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True) + # for i in range(uv_unposed_no_shape.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :]) + os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot'), exist_ok=True) + for i in range(uv_unposed_no_shape_one_hot.shape[0]): + np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape_one_hot[i, :, :, :]) # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True) # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body) diff --git a/convert_to_uv_local.py b/convert_to_uv_local.py new file mode 100644 index 0000000..3d2ede2 --- /dev/null +++ b/convert_to_uv_local.py @@ -0,0 +1,616 @@ +import os +import sys +import tqdm +import math +import time +import random +import shutil +import collections +import argparse +import threading +import numpy as np +import scipy as sp +import trimesh as tm +import cmapy as cm +import scipy.io as sio +import scipy.interpolate as spi +import matplotlib.pyplot as plt +import scipy.stats as stat +import cv2 +import plotly.express as ex +import plotly.graph_objects as go + + +from custom_utils.supersmpl_io import readPC2 + + + +os.environ["CUDA_VISIBLE_DEVICES"] = "2" + + +def add_triangle_centers(): + pass + +def dilate(img): + uv_mask_2 = img != 0 + m = img.min() + # norm = img.max() - img.min() + img = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.dilate(img, element, iterations=1) + result = (result + m) * uv_mask + return result + +def erode(img): + # m = img.min() + # norm = img.max() - img.min() + # img = (img - m) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.erode(img, element, iterations=2) + # result = result + m + return result + +def chunks(lst, n): + """Yield successive n-sized chunks from lst.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + + +class Result(): + def __init__(self) -> None: + self.result = dict() + +result_g = Result() + +class RBFThread(threading.Thread): + def __init__(self, uv_values, id, result): + threading.Thread.__init__(self) + self.uv_values = uv_values + self.id = id + self.result = result + + def run(self): + print('processing: ' + str(self.id)) + rbf = spi.RBFInterpolator(y=uv_map_discreet_unique, d=self.uv_values, neighbors=10) + res = rbf(rbf_dest) + self.result[self.id] = res + pass + +def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None): + verts_uv = np.zeros((res * 2, res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts + uv_values[m2uv[:, 1]] += verts + uv_values = uv_values[uv_map_mask, :] # add garment mask + uv_values /= uv_counter[:, np.newaxis] + return_garment_uv_mask = False + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + garment_mask = np.sum(uv_values, axis=-1) != 0 + + # test + verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # garment_mask = np.sum(uv_values, axis=-1) != 0 + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + rbf_dest_current = rbf_dest + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4) + garment_uv_mask_border = garment_uv_mask.astype(bool) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask + garment_uv_mask = garment_uv_mask * uv_mask[:,:,0] + rbf_dest_current = rbf_dest * garment_uv_mask + return_garment_uv_mask = True + + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values) + else: + raise NotImplementedError + + final = interpolator(rbf_dest_current) + + result = np.zeros((res * 2, res, 3)) + + # not sure, shouldn't be needed + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final + result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + if garment_uv_mask is not None: + result = result * garment_uv_mask[:,:,np.newaxis] + + # x = verts_uv #+ eroded + # uv_mask_2 = verts_uv != 0 + # y = result + # fig = plt.figure(figsize=(15, 15)) + # ax = fig.add_subplot(121) + # # im = ((x - x.min()) / (x.max() - x.min())) * uv_mask_2 + # im = ((x + 2) * 64 / 255) * uv_mask_2 + # ax.imshow(im) + # ax = fig.add_subplot(122) + # # im = ((y - y.min()) / (y.max() - y.min())) * uv_mask + # im = ((y + 2) * 64 / 255) * uv_mask # * (rbf_dest + garment_uv_mask)[:,:,np.newaxis] + # ax.imshow(im) + # plt.show() + + if return_garment_uv_mask: + return result, garment_uv_mask, garment_uv_mask_border + else: + return result + + + +def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = verts.shape[0] + verts_uv = np.zeros((batch, res * 2, res, 3)) + uv_values = np.zeros((batch, uv_map.shape[0], 3)) + uv_values[:, m2uv[:, 0]] = verts + uv_values[:, m2uv[:, 1]] += verts + uv_values = uv_values[:, uv_map_mask, :] + uv_values /= uv_counter[np.newaxis, :, np.newaxis] + return_garment_uv_mask = False + # test + verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0 + + unique_uv_values = uv_values[:, garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + rbf_dest_current = rbf_dest + + # add back mesh vertices that are not part of the garment + uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :] + rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) + rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + if garment_uv_mask is not None: + rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + result = np.zeros((batch, res * 2, res, 3), np.float32) + for i in range(batch): + print('processing frame: ' + str(i)) + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:]) + else: + raise NotImplementedError + result[i, rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current) + + result[:, uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values + + if garment_uv_mask is not None: + g_uv_mask = np.tile(garment_uv_mask[np.newaxis,:,:,np.newaxis], [batch,1,1,3]).astype(np.float32) + + result *= g_uv_mask + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + +def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = values.shape[0] + values = (values + 1) * mask + values_uv = np.zeros((batch, res * 2, res)) + missing = np.zeros((res * 2, res)) + uv_values = np.zeros((batch, uv_map.shape[0])) + uv_values[:, m2uv[:, 0]] = values + uv_values[:, m2uv[:, 1]] = values + uv_values = uv_values[:, uv_map_mask] + return_garment_uv_mask = False + # test + # verts_uv_save = verts_uv.copy() + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + raise Exception + + # uv_values = uv_values * garment_mask + + # test + values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0) + + unique_uv_values = uv_values[:, garment_mask] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0] + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:]) + final.append(interpolator(rbf_dest_current)) + final = np.asarray(final) + + result = np.zeros((batch, res * 2, res)) + result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final + if garment_uv_mask is not None: + result = result * garment_uv_mask[np.newaxis,:,:] + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + + +def m2uv_array(m2uv): + result = np.zeros((14475, 2), dtype=np.int32) + for k, v in m2uv.items(): + if len(v) == 2: + result[k] = list(v) + elif len(v) == 1: + x = v.pop() + result[k] = [x, x] + else: + assert False + + return result + + +def get_uv2m(m2uv): + + pass + +def uv_to_mesh(uv, m2uv): + res = uv.shape[0] + mesh = np.zeros((14475, 3)) + uv_pixels = uv_map_discreet[m2uv] + # counter = np.ones((14475)) + # counter[m2uv[:, 0] == m2uv[:, 1]] = 2 + mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis] + return mesh + pass + + +def verify_reconstruction(orig, reconstructed): + error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1)) + colors = np.array([cm.color('jet', 1 - i) for i in error]) + error_sum = np.sum(error) + + + verts_uv = np.zeros((uv_res * 2, uv_res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = colors + uv_values[m2uv[:, 1]] += colors + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(111) + ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255) + plt.show() + + + + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(121, projection='3d') + x = orig + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) + + ax = fig.add_subplot(122, projection='3d') + x = reconstructed#[error > 0.1, :] + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :] + + plt.show() + + pass + +# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs +def unique_pairs(a): + a.sort(axis=1) + b = np.ascontiguousarray(a).view( + np.dtype((np.void, a.dtype.itemsize * a.shape[1])) + ) + _, inv, ct = np.unique(b, return_inverse=True, return_counts=True) + return ct[inv] == 1 + + +def unique_uv_map_mask(uv_map_discreet): + mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool) + temp = np.zeros((uv_res * 2, uv_res)) + for p in uv_map_discreet: + temp[p[0], p[1]] += 1 + # temp = temp == 1 + for i, p in enumerate(uv_map_discreet): + if temp[p[0], p[1]] == 1: + mask[i] = True + elif temp[p[0], p[1]] == 2: + mask[i] = True + temp[p[0], p[1]] = 0 + return mask + + +def test(verts, faces, values): + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = values + uv_values[m2uv[:, 1]] += values + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + x = uv_values + verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) + im = ((x - x.min()) / (x.max() - x.min())) * 255 + tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show() + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train') + parser.add_argument("-s", "--start", help = "start of range", default=0) + parser.add_argument("-e", "--end", help = "end of range", default=10) + parser.add_argument("-n", "--name", help = "end of range", default='lhunor') + parser.add_argument("-sk", "--skirt", help = "end of range", default=0) + args = parser.parse_args() + + # smpl = SMPL('FeatureLines/superSMPL') + + local_folder = '/home/workspace/local-data/train' + + print("skirt" if args.skirt == 1 else "not skirt") + + # F = np.load("D:\\Work\\Code\\featurelines\\FeatureLines\\superSMPL\\smpl_faces.npy") + F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy") + # data = np.load("Blender/uv_data_1.npz", allow_pickle=True) + # uv_faces = data.f.Ft + # folder = 'Z:\\Work\\Huni\\Barcelona\\Dataset\\train' + folder = args.folder + uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy') + m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item() + m2uv = m2uv_array(m2uv) + uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy') + curvature_label_type = '6_25' + + uv_res = 256 + uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] + # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis] + # uv_mask_border = uv_mask * np.invert(uv_mask_dilate) + uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis] + + uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32) + uv_map_mask = unique_uv_map_mask(uv_map_discreet) + + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_counter = uv_counter[uv_map_mask] + + uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :] + + rbf_dest = np.ones((uv_res * 2, uv_res)) + rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0 + rbf_dest[np.invert(uv_mask[:,:,0])] = 0 + # rbf_dest = np.nonzero(rbf_dest) + # rbf_dest = np.array(rbf_dest).T + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + # rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order + + ids = os.listdir(folder)[int(args.start):int(args.end)] + + for id in tqdm.tqdm(ids): + garment_dir = os.path.join(folder, id) + + if not os.path.isdir(garment_dir): + continue + + # if os.path.exists(os.path.join(folder, id, 'uv_posed_body_no_orient')): + # continue + + 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 + if len (garments) != 1: + continue + garment = garments[0] + + # if os.path.exists(os.path.join(folder, id, 'uv_unposed', garment + '_0.npz')): + # continue + + # uv_verts_static = np.load(os.path.join(folder, id, garment + '_offsets_static_uv.npy')) + # uv_verts_static = ((uv_verts_static - uv_verts_static.min()) / (uv_verts_static.max() - uv_verts_static.min())) + # plt.imshow(uv_verts_static) + # plt.show() + # pass + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + frames = info['poses'].shape[1] + + # offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + verts_static = np.load(os.path.join(local_folder, 'mesh_static', id + '.npz'))['mesh_static'] + mask = np.load(os.path.join(local_folder, 'mesh_mask', id + '.npy')).astype(bool) + # body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] + # gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + # curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy')) + # mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])]) + # mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])]) + # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])]) + mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz')) for i in range(frames)] + mesh_unposed_no_shape_one_hot = np.stack([mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)]) + + + # verts_dynamic = gt + + # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing? + + # test(data.f.Vt, data.f.Ft, verts) + # test(uv_map, data.f.Ft, verts) + + start = time.perf_counter() + + + + # fig = go.Figure(data=[ + # go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,1], + # z=x[:,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=y[:,0], + # y=y[:,1], + # z=y[:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2], + # ), + # ]) + + # fig.update_layout(scene_aspectmode='manual', + # scene_aspectratio=dict(x=1, y=1, z=1)) + # fig.update_layout( + # scene = dict( + # xaxis = dict(nticks=4, range=[-1,1],), + # yaxis = dict(nticks=4, range=[-1,1],), + # zaxis = dict(nticks=4, range=[-1,1],),) + # ) + + # fig.show() + + uv_verts_static, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, verts_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None) + uv_verts_static = uv_verts_static.astype(np.float32) + # uv_body_static = convert_to_uv(uv_res, m2uv, uv_map, body_static, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) + # uv_body = convert_to_uv_batch(uv_res, m2uv, uv_map, body[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_posed_body_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, body_no_orient[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) + print(time.perf_counter() - start) + + start = time.perf_counter() + # np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static) + # np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static) + # np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask) + + # os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True) + # for i in range(uv_body.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'body_posed_uv', garment + '_' + str(i) + '.npz'), uv_body=uv_body[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_posed_body_no_orient'), exist_ok=True) + # for i in range(uv_posed_body_no_orient.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True) + # for i in range(uv_gt.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True) + # for i in range(uv_curvature.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True) + # for i in range(uv_unposed.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True) + # for i in range(uv_unposed_no_shape.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :]) + + os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot'), exist_ok=True) + for i in range(uv_unposed_no_shape_one_hot.shape[0]): + np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape_one_hot[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True) + # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body) + + print(time.perf_counter() - start) + + # mesh = uv_to_mesh(uv_verts_static, m2uv) + # verify_reconstruction(verts_static * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + # mesh = uv_to_mesh(uv_body[0,:,:], m2uv) + # verify_reconstruction(body[0,:,:] * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + pass + diff --git a/convert_to_uv_local_single.py b/convert_to_uv_local_single.py new file mode 100644 index 0000000..06b9867 --- /dev/null +++ b/convert_to_uv_local_single.py @@ -0,0 +1,539 @@ +import os +import sys +import glob +import tqdm +import math +import time +import random +import shutil +import collections +import argparse +import threading +import numpy as np +import scipy as sp +import trimesh as tm +import cmapy as cm +import scipy.io as sio +import scipy.interpolate as spi +import matplotlib.pyplot as plt +import scipy.stats as stat +import cv2 +import plotly.express as ex +import plotly.graph_objects as go + + +os.environ["CUDA_VISIBLE_DEVICES"] = "2" + + +def add_triangle_centers(): + pass + +def dilate(img): + uv_mask_2 = img != 0 + m = img.min() + # norm = img.max() - img.min() + img = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.dilate(img, element, iterations=1) + result = (result + m) * uv_mask + return result + +def erode(img): + # m = img.min() + # norm = img.max() - img.min() + # img = (img - m) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.erode(img, element, iterations=2) + # result = result + m + return result + +def chunks(lst, n): + """Yield successive n-sized chunks from lst.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + + +def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None): + verts_uv = np.zeros((res * 2, res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts + uv_values[m2uv[:, 1]] += verts + uv_values = uv_values[uv_map_mask, :] # add garment mask + uv_values /= uv_counter[:, np.newaxis] + return_garment_uv_mask = False + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + garment_mask = np.sum(uv_values, axis=-1) != 0 + + # test + verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # garment_mask = np.sum(uv_values, axis=-1) != 0 + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + rbf_dest_current = rbf_dest + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4) + garment_uv_mask_border = garment_uv_mask.astype(bool) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask + garment_uv_mask = garment_uv_mask * uv_mask[:,:,0] + rbf_dest_current = rbf_dest * garment_uv_mask + return_garment_uv_mask = True + + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values) + else: + raise NotImplementedError + + final = interpolator(rbf_dest_current) + + result = np.zeros((res * 2, res, 3)) + + # not sure, shouldn't be needed + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final + result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + if garment_uv_mask is not None: + result = result * garment_uv_mask[:,:,np.newaxis] + + # x = verts_uv #+ eroded + # uv_mask_2 = verts_uv != 0 + # y = result + # fig = plt.figure(figsize=(15, 15)) + # ax = fig.add_subplot(121) + # # im = ((x - x.min()) / (x.max() - x.min())) * uv_mask_2 + # im = ((x + 2) * 64 / 255) * uv_mask_2 + # ax.imshow(im) + # ax = fig.add_subplot(122) + # # im = ((y - y.min()) / (y.max() - y.min())) * uv_mask + # im = ((y + 2) * 64 / 255) * uv_mask # * (rbf_dest + garment_uv_mask)[:,:,np.newaxis] + # ax.imshow(im) + # plt.show() + + if return_garment_uv_mask: + return result, garment_uv_mask, garment_uv_mask_border + else: + return result + + + +def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None): + batch = len(verts) + # verts_uv = np.zeros((batch, res * 2, res, 3)) + + # test + # verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + assert method == 'rbf' + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + assert mask is not None + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + + rbf_dest_current = rbf_dest + + # add back mesh vertices that are not part of the garment + uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :] + rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) + rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True + + assert not (garment_uv_mask is None and mask is not None) + assert garment_uv_mask is not None + rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + + for i in range(batch): + print('processing frame: ' + str(i)) + + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts[i] + uv_values[m2uv[:, 1]] += verts[i] + uv_values = uv_values[uv_map_mask, :] + uv_values /= uv_counter[:, np.newaxis] + + + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + + + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear') + result = np.zeros((res * 2, res, 3), np.float32) + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current) + result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:] + + if garment_uv_mask is not None: + result *= garment_uv_mask[:,:,np.newaxis] + + np.savez_compressed(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2', id + '_' + str(i) + '.npz'), uv_unposed_no_shape=result) + + + +def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = values.shape[0] + values = (values + 1) * mask + values_uv = np.zeros((batch, res * 2, res)) + missing = np.zeros((res * 2, res)) + uv_values = np.zeros((batch, uv_map.shape[0])) + uv_values[:, m2uv[:, 0]] = values + uv_values[:, m2uv[:, 1]] = values + uv_values = uv_values[:, uv_map_mask] + return_garment_uv_mask = False + # test + # verts_uv_save = verts_uv.copy() + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + raise Exception + + # uv_values = uv_values * garment_mask + + # test + values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0) + + unique_uv_values = uv_values[:, garment_mask] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0] + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:]) + final.append(interpolator(rbf_dest_current)) + final = np.asarray(final) + + result = np.zeros((batch, res * 2, res)) + result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final + if garment_uv_mask is not None: + result = result * garment_uv_mask[np.newaxis,:,:] + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + + +def m2uv_array(m2uv): + result = np.zeros((14475, 2), dtype=np.int32) + for k, v in m2uv.items(): + if len(v) == 2: + result[k] = list(v) + elif len(v) == 1: + x = v.pop() + result[k] = [x, x] + else: + assert False + + return result + + +def get_uv2m(m2uv): + + pass + +def uv_to_mesh(uv, m2uv): + res = uv.shape[0] + mesh = np.zeros((14475, 3)) + uv_pixels = uv_map_discreet[m2uv] + # counter = np.ones((14475)) + # counter[m2uv[:, 0] == m2uv[:, 1]] = 2 + mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis] + return mesh + pass + + +def verify_reconstruction(orig, reconstructed): + error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1)) + colors = np.array([cm.color('jet', 1 - i) for i in error]) + error_sum = np.sum(error) + + + verts_uv = np.zeros((uv_res * 2, uv_res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = colors + uv_values[m2uv[:, 1]] += colors + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(111) + ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255) + plt.show() + + + + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(121, projection='3d') + x = orig + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) + + ax = fig.add_subplot(122, projection='3d') + x = reconstructed#[error > 0.1, :] + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :] + + plt.show() + + pass + +# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs +def unique_pairs(a): + a.sort(axis=1) + b = np.ascontiguousarray(a).view( + np.dtype((np.void, a.dtype.itemsize * a.shape[1])) + ) + _, inv, ct = np.unique(b, return_inverse=True, return_counts=True) + return ct[inv] == 1 + + +def unique_uv_map_mask(uv_map_discreet): + mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool) + temp = np.zeros((uv_res * 2, uv_res)) + for p in uv_map_discreet: + temp[p[0], p[1]] += 1 + # temp = temp == 1 + for i, p in enumerate(uv_map_discreet): + if temp[p[0], p[1]] == 1: + mask[i] = True + elif temp[p[0], p[1]] == 2: + mask[i] = True + temp[p[0], p[1]] = 0 + return mask + + +def test(verts, faces, values): + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = values + uv_values[m2uv[:, 1]] += values + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + x = uv_values + verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) + im = ((x - x.min()) / (x.max() - x.min())) * 255 + tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show() + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/train') + parser.add_argument("-s", "--start", help = "start of range", default=0) + parser.add_argument("-e", "--end", help = "end of range", default=10) + parser.add_argument("-n", "--name", help = "end of range", default='lhunor') + parser.add_argument("-sk", "--skirt", help = "end of range", default=0) + args = parser.parse_args() + + # smpl = SMPL('FeatureLines/superSMPL') + + local_folder = args.folder + os.makedirs(os.path.join(local_folder, 'uv_unposed_no_shape_one_hot_2'), exist_ok=True) + + print("skirt" if args.skirt == 1 else "not skirt") + + # F = np.load("D:\\Work\\Code\\featurelines\\FeatureLines\\superSMPL\\smpl_faces.npy") + F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy") + # data = np.load("Blender/uv_data_1.npz", allow_pickle=True) + # uv_faces = data.f.Ft + # folder = 'Z:\\Work\\Huni\\Barcelona\\Dataset\\train' + uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy') + m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item() + m2uv = m2uv_array(m2uv) + uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy') + curvature_label_type = '6_25' + + uv_res = 256 + uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] + # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis] + # uv_mask_border = uv_mask * np.invert(uv_mask_dilate) + uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis] + + uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32) + uv_map_mask = unique_uv_map_mask(uv_map_discreet) + + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_counter = uv_counter[uv_map_mask] + + uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :] + + rbf_dest = np.ones((uv_res * 2, uv_res)) + rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0 + rbf_dest[np.invert(uv_mask[:,:,0])] = 0 + # rbf_dest = np.nonzero(rbf_dest) + # rbf_dest = np.array(rbf_dest).T + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + # rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order + + ids = os.listdir(os.path.join(local_folder, 'mesh_static'))[int(args.start):int(args.end)] + ids = [id.split('.')[0] for id in ids] + + for id in tqdm.tqdm(ids): + + + # if os.path.exists(os.path.join(folder, id, 'uv_unposed', garment + '_0.npz')): + # continue + + # uv_verts_static = np.load(os.path.join(folder, id, garment + '_offsets_static_uv.npy')) + # uv_verts_static = ((uv_verts_static - uv_verts_static.min()) / (uv_verts_static.max() - uv_verts_static.min())) + # plt.imshow(uv_verts_static) + # plt.show() + # pass + + frames = glob.glob(os.path.join(local_folder, 'pose_shape', id + '_*')) + frames = len(frames) + assert frames <= 300 + + # offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + verts_static = np.load(os.path.join(local_folder, 'mesh_static', id + '.npz'))['mesh_static'] + mask = np.load(os.path.join(local_folder, 'mesh_mask', id + '.npy')).astype(bool) + # body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] + # gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + # curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy')) + # mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])]) + # mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])]) + # body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])] + # body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])]) + mesh_unposed_no_shape_one_hot = [np.load(os.path.join(local_folder, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz')) for i in range(frames)] + mesh_unposed_no_shape_one_hot = [mesh_unposed_no_shape_one_hot[i]['mesh_unposed_no_shape'] for i in range(frames)] + + + # verts_dynamic = gt + + # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing? + + # test(data.f.Vt, data.f.Ft, verts) + # test(uv_map, data.f.Ft, verts) + + start = time.perf_counter() + + + + uv_verts_static, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, verts_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None) + uv_verts_static = uv_verts_static.astype(np.float32) + # uv_body_static = convert_to_uv(uv_res, m2uv, uv_map, body_static, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) + # uv_body = convert_to_uv_batch(uv_res, m2uv, uv_map, body[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_posed_body_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, body_no_orient[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + # uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32) + uv_unposed_no_shape_one_hot = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape_one_hot, neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border, id=id) + # uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) + print(time.perf_counter() - start) + + # start = time.perf_counter() + # np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static) + # np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static) + # np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask) + + # os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True) + # for i in range(uv_body.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'body_posed_uv', garment + '_' + str(i) + '.npz'), uv_body=uv_body[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_posed_body_no_orient'), exist_ok=True) + # for i in range(uv_posed_body_no_orient.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True) + # for i in range(uv_gt.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True) + # for i in range(uv_curvature.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True) + # for i in range(uv_unposed.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True) + # for i in range(uv_unposed_no_shape.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot'), exist_ok=True) + # for i in range(uv_unposed_no_shape_one_hot.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape_one_hot[i, :, :, :]) + + # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True) + # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body) + + # print(time.perf_counter() - start) + + # mesh = uv_to_mesh(uv_verts_static, m2uv) + # verify_reconstruction(verts_static * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + # mesh = uv_to_mesh(uv_body[0,:,:], m2uv) + # verify_reconstruction(body[0,:,:] * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + pass +