import os import sys import tqdm import math import time import random import shutil import argparse import numpy as np import scipy as sp import scipy.io as sio import cv2 from custom_utils.supersmpl_io import readPC2 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 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 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 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 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() return error_sum pass if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-fi", "--folder_in", help = "end of range") parser.add_argument("-fo", "--folder_out", help = "end of range") parser.add_argument("-s", "--start", help = "start of range") parser.add_argument("-e", "--end", help = "end of range") args = parser.parse_args() folder = args.folder_in output_folder = args.folder_out # # F = np.load("D:\\Work\\Code\\featurelines\\FeatureLines\\superSMPL\\smpl_faces.npy") # F = np.load("FeatureLines/superSMPL/smpl_faces.npy") # # data = np.load("Blender/uv_data_1.npz", allow_pickle=True) # # uv_faces = data.f.Ft # # folder = 'D:\\Work\\Dataset\\validation_2' # uv_map = np.load('MeshUtils/stretched_uv_map_3.npy') # m2uv = np.load('MeshUtils/m2uv.npy', allow_pickle=True).item() # m2uv = m2uv_array(m2uv) # uv_mask = np.load('MeshUtils/mask_256x512.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 ids = os.listdir(folder)[int(args.start):int(args.end)] os.makedirs(os.path.join(output_folder), exist_ok=True) # os.makedirs(os.path.join(output_folder, 'uv_shaped_body'), exist_ok=True) 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, 'gt_uv')): # 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 garment != 'Jumpsuit': # continue start = time.perf_counter() # 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)) # verts_static = (body_static + offsets_static) * mask[:, np.newaxis] # dfassert verts_static.min() > -1 and verts_static.max() < 1 # os.makedirs(os.path.join(output_folder, 'mesh_static'), exist_ok=True) # np.savez_compressed(os.path.join(output_folder, 'mesh_static', id + '.npz'), mesh_static=verts_static) # os.makedirs(os.path.join(output_folder, 'uv_static'), exist_ok=True) # uv_verts_static = np.load(os.path.join(folder, id, garment + '_offsets_static_uv.npy')) # np.savez_compressed(os.path.join(output_folder, 'uv_static', id + '.npz'), uv_static=uv_verts_static) # garment_folder_gt = os.path.join(folder, id, 'gt_uv') # garment_folder_body = os.path.join(folder, id, 'body_posed_uv') # garment_folder_body_mesh = os.path.join(folder, id, 'mesh_body_posed') # os.makedirs(os.path.join(output_folder, 'mesh_mask'), exist_ok=True) # shutil.copyfile(os.path.join(folder, id, garment + '_mask.npy'), os.path.join(output_folder, 'mesh_mask', id + '.npy')) # os.makedirs(os.path.join(output_folder, 'mesh_posed'), exist_ok=True) # garment_posed = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V'] # for i in range(garment_posed.shape[0]): # np.savez_compressed(os.path.join(output_folder, 'mesh_posed', id + '_' + str(i) + '.npz'), mesh_posed=garment_posed[i,:,:] * mask[:, np.newaxis]) # os.makedirs(os.path.join(output_folder, 'mesh_body_posed'), exist_ok=True) body_posed = readPC2(os.path.join(garment_dir, 'Body_posed.pc16'), True)['V'] # for i in range(body_posed.shape[0]): # np.save(os.path.join(output_folder, 'mesh_body_posed', id + '_' + str(i) + '.npy'), body_posed[i,:,:]) # os.makedirs(os.path.join(output_folder, 'uv_posed'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_folder_gt, garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_posed', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'uv_body_posed'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_folder_body, garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_body_posed', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'pose_shape'), exist_ok=True) # info = sio.loadmat(os.path.join(garment_dir, 'info.mat')) # pose, shape, gender = info['poses'], info['shape'], info['gender'] # shape = np.squeeze(np.concatenate((shape, gender), axis=1)) # for i in range(body_posed.shape[0]): # np.save(os.path.join(output_folder, 'pose_shape', id + '_' + str(i) + '.npy'), np.concatenate((pose[:, i], shape), axis=-1)) # os.makedirs(os.path.join(output_folder, 'uv_mask'), exist_ok=True) # assert os.path.exists(os.path.join(folder, id, 'mask_uv.npy')) # if not os.path.exists(os.path.join(output_folder, 'uv_mask', id + '.npz')): # uv_mask = np.load(os.path.join(folder, id, 'mask_uv.npy')) # np.savez_compressed(os.path.join(output_folder, 'uv_mask', id + '.npz'), uv_mask=uv_mask) # os.makedirs(os.path.join(output_folder, 'mesh_unposed'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'mesh_unposed', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_unposed', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'uv_unposed'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'uv_unposed', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_unposed', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'mesh_unposed_no_shape'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'uv_unposed_no_shape'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_unposed_no_shape', id + '_' + str(i) + '.npz')) os.makedirs(os.path.join(output_folder, 'uv_curvature'), exist_ok=True) for i in range(body_posed.shape[0]): shutil.copyfile(os.path.join(garment_dir, 'curvature_uv', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_curvature', id + '_' + str(i) + '.npz')) os.makedirs(os.path.join(output_folder, 'mesh_curvature'), exist_ok=True) mesh_curvature = np.load(os.path.join(garment_dir, 'curvature_label_6_25.npy')) for i in range(body_posed.shape[0]): np.savez_compressed(os.path.join(output_folder, 'mesh_curvature', id + '_' + str(i) + '.npz'), mesh_curvature=mesh_curvature[i]) # shutil.copyfile(os.path.join(folder, id, 'uv_shaped_body.npz'), os.path.join(output_folder, 'uv_shaped_body', id + '.npz')) # os.makedirs(os.path.join(output_folder, 'uv_posed_body_no_orient'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_posed_body_no_orient', id + '_' + str(i) + '.npz')) # os.makedirs(os.path.join(output_folder, 'mesh_posed_body_no_orient'), exist_ok=True) # for i in range(body_posed.shape[0]): # shutil.copyfile(os.path.join(garment_dir, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) pass