luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/convert_to_uv.py
+5 -1
View File
@@ -399,10 +399,11 @@ def get_shapedirs():
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
@@ -442,6 +443,7 @@ if __name__ == '__main__':
#########################################################################################
folder = args.folder
# data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed'] # 'mesh_unposed',
data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
@@ -460,6 +462,8 @@ if __name__ == '__main__':
os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
# ids = ['00438']
for id in tqdm.tqdm(ids):
garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
data/convert_to_uv_cape.py
+33 -34
View File
@@ -401,7 +401,7 @@ def get_shapedirs():
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -448,11 +448,11 @@ if __name__ == '__main__':
folder = args.folder
# data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed',
data_to_convert = ['mesh_posed', 'mesh_static'] # 'mesh_unposed',
frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end]
frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))
frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames]
ids = list(set([frame[0] for frame in frames]))
ids = sorted(list(set([frame[0] for frame in frames])))[args.start:args.end]
frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids}
shapedirs = get_shapedirs()
@@ -462,11 +462,11 @@ if __name__ == '__main__':
os.makedirs(folder, exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
# os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
mask = np.ones((3638), dtype=bool)
@@ -476,15 +476,14 @@ if __name__ == '__main__':
# x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
# mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
# mesh_posed[:,mask, :] = x
if 'mesh_unposed' in data_to_convert:
mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
# if 'mesh_unposed' in data_to_convert:
# mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
# mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
if 'mesh_static' in data_to_convert:
mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
if 'mesh_normals' in data_to_convert:
mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0)
if 'body_posed' in data_to_convert and k != 1:
body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
if 'mesh_posed' in data_to_convert:
mesh_posed = np.stack([np.load(os.path.join(folder, 'mesh_posed_no_orient', f'{id}.{int(i):06d}.npz'))['mesh_posed_no_orient'] for i in frames_by_id[id]], axis=0)
start = time.perf_counter()
@@ -493,13 +492,13 @@ if __name__ == '__main__':
uv_static = uv_static[0].astype(np.float32)
if 'mesh_posed' in data_to_convert:
uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_unposed' in data_to_convert:
uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'mesh_normals' in data_to_convert:
uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
if 'body_posed' in data_to_convert and k != 1:
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
uv_posed_no_orient = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'mesh_unposed' in data_to_convert:
# uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'mesh_normals' in data_to_convert:
# uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
# if 'body_posed' in data_to_convert and k != 1:
# uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
print(time.perf_counter() - start)
@@ -515,16 +514,16 @@ if __name__ == '__main__':
# np.savez_compressed(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
if 'mesh_posed' in data_to_convert:
for i in range(len(uv_posed)):
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
if 'mesh_unposed' in data_to_convert:
for i in range(len(uv_unposed)):
np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
if 'mesh_normals' in data_to_convert:
for i, frame_nr in enumerate(frames_by_id[id]):
np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
if 'body_posed' in data_to_convert and k != 1:
for i in range(len(uv_body_posed)):
np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
for k, i in enumerate(frames_by_id[id]):
np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient', f'{id}.{int(i):06d}.npz'), uv_posed_no_orient=uv_posed_no_orient[k])
# if 'mesh_unposed' in data_to_convert:
# for i in range(len(uv_unposed)):
# np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
# if 'mesh_normals' in data_to_convert:
# for i, frame_nr in enumerate(frames_by_id[id]):
# np.savez_compressed(os.path.join(folder, 'uv_normals', f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
# if 'body_posed' in data_to_convert and k != 1:
# for i in range(len(uv_body_posed)):
# np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient', f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
# pass
data/dataset.py
+11 -1
View File
@@ -276,7 +276,7 @@ if __name__ == '__main__':
'uv_mask',
}
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
dataset = Dataset('/data-local/127_2/cloth3d_processed_test', signature, resolutions=[512], device=device, frames_file='frames_temp.txt')
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
smpl = SMPL(orig_smpl=False, path='smpl', device='cuda')
@@ -293,6 +293,16 @@ if __name__ == '__main__':
step = 0
for batch in train_dataloader:
batch
from custom_utils.custom_utils import save_stl
fix_data = np.load(os.path.join(dataset.folder, 'mesh_fixed', f'{batch["id"][0]}_{batch["garment"][0]}.npz'))
faces = fix_data['faces']
mask = fix_data['mask']
save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_posed.stl', batch['mesh_posed'][0].cpu().numpy()[mask], faces)
save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_unposed.stl', batch['mesh_unposed'][0].cpu().numpy()[mask], faces)
print(batch['id'][0])
with profiler:
m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]
data/downscale.py
+5 -23
View File
@@ -11,7 +11,7 @@ import plotly.express as ex
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.utils import Frame
@@ -165,32 +165,19 @@ def downscale_uv_maps_static(folder, type, sizes):
np.savez_compressed(os.path.join(folder, type + '_' + str(size), f'{frame[0]}_{frame[1]}.npz'), uv_static = down)
def downscale_uv_maps_mask(folder, type, size):
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = sorted(os.listdir(os.path.join(folder, type)))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
id, _ = frame.split('.')
img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
down = np.squeeze(down.numpy().astype(bool))
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cloth3d_processed_test')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=25000)
parser.add_argument("-e", "--end", help = "end of range", default=300000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
args = parser.parse_args()
# sizes = [32, 64, 128, 256]
device = 'cuda'
sizes = [32, 64, 128, 256]
sizes = [64, 128, 256]
folder = args.folder
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
@@ -199,19 +186,14 @@ if __name__ == '__main__':
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_masks_body[512] = uv_mask_body.astype(bool)
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:].astype(np.int32)).to(device))
down = np.squeeze(down.cpu().numpy().astype(bool))
uv_masks_body[res] = down
uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
uv_mask_store[512] = UVMaskStore(folder, size=512)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
data/downscale_cape.py
+27 -7
View File
@@ -23,7 +23,7 @@ def get_img_idxs():
img_idxs[512] = img_idx
img_coords[512] = coords
uv_masks[512] = uv_mask.astype(np.int32)
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy()
down = np.squeeze(down.astype(bool))
coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
@@ -37,7 +37,7 @@ def get_img_idxs():
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_masks_body[512] = uv_mask_body
for res in [32, 64, 128, 256]:
for res in [64, 128, 256]:
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy()
down = np.squeeze(down.astype(bool))
uv_masks_body[res] = down
@@ -113,6 +113,8 @@ def downscale_uv_normals(folder, type, sizes):
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
if not frame.startswith('03223_shortshort_chicken_wings'):
continue
img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals']
img = fill(img, uv_masks_body[512])
for size in sizes:
@@ -140,6 +142,23 @@ def downscale_uv_maps_unposed(folder, type, sizes):
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
def downscale_uv_maps_posed(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames):
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_no_orient']
mask_full = uv_mask_store[512][None]
img = fill(img, mask_full)
for size in sizes:
mask = uv_mask_store[size][None]
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
down = down * mask[None,:,:]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_no_orient = down)
def downscale_uv_maps_static(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
@@ -173,7 +192,7 @@ def downscale_uv_maps_mask(folder, type, size):
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male')
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/female')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=100000)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
@@ -182,7 +201,7 @@ if __name__ == '__main__':
device = 'cuda'
sizes = [32, 64, 128, 256]
sizes = [64, 128, 256]
folder = args.folder
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
@@ -199,9 +218,10 @@ if __name__ == '__main__':
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
# downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
downscale_uv_maps_posed(folder=folder, type='uv_posed_no_orient', sizes=sizes)
# downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
# downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
print()
data/preprocessing.py
+184 -10
View File
@@ -12,6 +12,7 @@ args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import math
import shutil
@@ -76,6 +77,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
ids = ['00021']
for id in tqdm(ids):
# id = '01467'
@@ -87,8 +90,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
info = loadInfo(os.path.join(folder, id, 'info.mat'))
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
# if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
# continue
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
# if info['poses'].shape[1] == verts.shape[0]:
@@ -488,6 +491,176 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
pass
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda')
from scipy.spatial import KDTree
from pytorch3d.structures import Meshes
from torch_cluster import knn
def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count)
batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475)
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1]
idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
return idx
def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004):
F_norm_local = F_norm
mesh = Meshes(B, F_norm_local)
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
idx = _nearest_neighbour(V, B)
# compute loss
D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
dot = torch.einsum('abc,abc->ab', D, b_vn)
thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
loss = torch.sum(thresholded ** 2, axis=-1)
return loss
def area_loss(pred, gt, faces):
"""
Calculates the area loss between predicted and ground truth meshes.
Args:
pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3).
gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3).
faces (torch.Tensor): Mesh faces of shape (num_faces, 3).
Returns:
torch.Tensor: Area loss between predicted and ground truth meshes.
"""
# Calculate the area of each face for predicted and ground truth meshes
pred_faces = pred[:, faces, :]
gt_faces = gt[:, faces, :]
pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :]
pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :]
gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :]
gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :]
pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2)
gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2)
pred_area = 0.5 * torch.norm(pred_cross, dim=2)
gt_area = 0.5 * torch.norm(gt_cross, dim=2)
# Calculate the area loss
area_loss = torch.mean(torch.abs(pred_area - gt_area))
return area_loss
def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1):
ids = sorted(os.listdir(folder))[start:end]
os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True)
F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
ids = ['00005']
for id in tqdm(ids):
garment_dir = os.path.join(folder, id)
if not os.path.isdir(garment_dir):
continue
garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:]
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
smpl = SMPL(path='smpl', device='cuda', batch_size=1)
pose = np.zeros((24, 3), np.float32)
pose[1,2] = .15
pose[2,2] = -.15
pose = torch.tensor(pose.flatten()[None,:], dtype=torch.float32, device='cuda')
for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue
mesh_static, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj'))
faces = quads2tris(faces)
mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
unposed = mesh_static
unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda')
unposed.requires_grad = True
mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda')
optimizer = torch.optim.Adam([unposed], lr=0.0001)
for i in range(10000):
optimizer.zero_grad()
body = smpl(pose, shape)[0]
c_loss = collision_loss(unposed, body)
a_loss = area_loss(unposed, mesh_static, faces)
loss : torch.Tensor = c_loss + a_loss
# Backward pass
loss.backward()
optimizer.step()
pose = pose / 1.0005
# Print the loss
if i % 50 == 0:
pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item()
loss_str = loss.detach().cpu().numpy()[0]
c_loss_str = c_loss.detach().cpu().numpy()[0]
a_loss_str = a_loss.detach().cpu().numpy()
print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}")
unposed = unposed.detach().cpu().numpy()
body = body.detach().cpu().numpy()
gender = int(shape[0].cpu().numpy()[-1] > 0)
shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
unshaped = (unposed - shape_offsets)
writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
verts = unposed
# verts = unshaped
frame = 0
fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# add body mesh
fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1))
fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube')
fig.show()
# verts = pred.detach().cpu().numpy()[:,mask,:]
# frame = 9
# fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
# fig.update_layout(scene = dict(
# xaxis = dict(range=[-1,1],),
# yaxis = dict(range=[-1,1],),
# zaxis = dict(range=[-1,1],),),)
# fig.update_layout(scene_aspectmode='cube')
# fig.show()
pass
def create_frames_file(folder, out_folder, split=0.9):
ids = os.listdir(folder)
@@ -879,13 +1052,13 @@ def generate_fixed_masks(folder, out_folder):
if __name__ == '__main__':
# folder = '/data-net/dataset/cloth3d/test'
# folder_reg = '/data-local/127/cloth3d_registration/test'
# out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/test'
folder_reg = '/data-local/127/cloth3d_registration/test'
out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/train'
folder_reg = '/data-local/127/cloth3d_registration/train'
out_folder = '/data-local/127/cloth3d_processed'
# folder = '/data-net/dataset/cloth3d/train'
# folder_reg = '/data-local/127/cloth3d_registration/train'
# out_folder = '/data-local/127/cloth3d_processed'
smpl = SMPL(path='smpl', device='cuda')
@@ -898,13 +1071,14 @@ if __name__ == '__main__':
# remove_irregular_poses(folder, folder_reg, out_folder)
# double_check_irregular_poses(folder, folder_reg, out_folder)
# create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
# verify_registration_ans_mask(folder, out_folder)
# create_pose_shape(folder, out_folder)
# create_mesh_static(folder, out_folder)
# create_mesh_body_posed_no_orient(folder, out_folder)
# create_mesh_normals(folder, out_folder)
# create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
# create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end)
# calculate_average(folder, out_folder)
# create_frames_file(folder, out_folder)
@@ -915,4 +1089,4 @@ if __name__ == '__main__':
# folder = '/data-local/127_2/cape_processed/test/female'
# create_mesh_normals_cape(folder, None)
generate_fixed_masks(folder, out_folder)
# generate_fixed_masks(folder, out_folder)
data/utils.py
+20 -4
View File
@@ -1,10 +1,26 @@
import tables as tb
class Frame:
def __init__(self, id: str, frame_nr: int, garment: str=''):
self.id: str = id
self.frame_nr: int = frame_nr
self.garment: str = garment
def __init__(self, id: str ='', frame_nr: int= 0, garment: str='', frame = None, frame_str: str = ''):
if frame is not None:
self.id: str = frame.id
self.frame_nr: int = frame.frame_nr
self.garment: str = frame.garment
assert id == '' and frame_nr == 0 and garment == '' and frame_str == ''
elif frame_str != '':
if ' ' in frame_str:
self.id: str = frame_str.split()[0]
self.frame_nr: int = int(frame_str.split()[1])
self.garment: str = frame_str.split()[2]
else:
self.id: str = frame_str.split('_')[0]
self.frame_nr: int = int(frame_str.split('_')[1])
self.garment: str = frame_str.split('_')[2]
assert id == '' and frame_nr == 0 and garment == '' and frame is None
else:
self.id: str = id
self.frame_nr: int = frame_nr
self.garment: str = garment
def __hash__(self):
return hash((self.id, self.frame_nr, self.garment))