278 lines
12 KiB
Python
278 lines
12 KiB
Python
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, 'mesh_unposed_no_shape_one_hot'), exist_ok=True)
|
|
# for i in range(pose.shape[1]):
|
|
# shutil.copyfile(os.path.join(garment_dir, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_unposed_no_shape_one_hot', 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_unposed_no_shape_one_hot'), exist_ok=True)
|
|
# for i in range(body_posed.shape[0]):
|
|
# shutil.copyfile(os.path.join(garment_dir, 'uv_unposed_no_shape_one_hot', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_unposed_no_shape_one_hot', 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(pose.shape[1]):
|
|
np.savez_compressed(os.path.join(output_folder, 'mesh_curvature', id + '_' + str(i) + '.npz'), mesh_curvature=mesh_curvature[i])
|
|
|
|
|
|
# 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_2', 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'))
|
|
|
|
# shutil.copyfile(os.path.join(folder, id, 'uv_shaped_body.npz'), os.path.join(output_folder, 'uv_shaped_body', id + '.npz'))
|
|
|
|
pass
|
|
|