import os import sys # os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' import math import time import torch import random import pytorch3d import torchvision as tv import numpy as np import scipy.io as sio import plotly.express as ex import torchvision from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.utils import Frame from data.cloth3d_io import readPC2, readPC2Frame from data.dataset_cape import DatasetCape def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): if frames_file is None: uv_static_folder = os.path.join(folder, 'pose_shape') frames = sorted(os.listdir(uv_static_folder)) if is_cape: frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames] else: frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] # if testing: # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) # frames = [frame.split('.')[0] + '_4' for frame in frames] if frames_file is not None: frames = [] with open(os.path.join(folder, frames_file), 'r') as f: lines = f.readlines() lines = [line.strip() for line in lines] frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines] if shuffle: random.shuffle(frames) return frames class StaticCache(): def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None: self.store = dict() self.signature = signature self.folder = folder assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) self.resolutions = resolutions self.is_cape = is_cape self.device = device if is_cape: self.uv_masks = dict() uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() for res in [32, 64, 128, 256]: down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:]) down = torch.squeeze(down).to(torch.float32) self.uv_masks[res] = down.cpu().numpy() def read(self, id, garment): data = dict() mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool) data['mesh_mask'] = mesh_mask for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_static' + res_string in self.signature: uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static if 'uv_mask' + res_string in self.signature: if self.is_cape: data['uv_mask' + res_string] = self.uv_masks[res] else: uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask if 'mesh_static' in self.signature: mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy')) data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32) data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32) # if 'uv_shaped_body' in self.signature: # uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body return data def __getitem__(self, frame: Frame) -> dict: if (frame.id, frame.garment) not in self.store: item = self.read(frame.id, frame.garment) stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} for key, value in stored_data.items(): if key.startswith('uv_') and value.shape[-1] == 3: stored_data[key] = torch.permute(value, [2,0,1]) self.store[(frame.id, frame.garment)] = stored_data data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()} return data def read_frames_file(frames_file, shuffle=False): frames = [] with open(frames_file, 'r') as f: lines = f.readlines() lines = [line.strip() for line in lines] frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines] if shuffle: random.shuffle(frames) return frames class DatasetCloth3D(torch.utils.data.Dataset): def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): self.folder = folder self.device = device self.is_cape = is_cape self.signature = signature self.resolutions = resolutions self.transforms = tv.transforms.ToTensor() self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=False) self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None ids = set([frame.id for frame in self.frames]) self.pose_cache = dict() for id in ids: self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy')) self.shape_cache = dict() for id in ids: self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy')) def __len__(self): return len(self.frames) def __getitem__(self, idx): frame : Frame = self.frames[idx] if self.cache is not None: if frame in self.cache.keys(): data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()} data = { **data, **self.static_cache[frame] } return data item = dict() mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool) item['id'] = frame.id item['frame_nr'] = frame.frame_nr item['garment'] = frame.garment if 'pose_shape' in self.signature: pose = self.pose_cache[frame.id][frame.frame_nr] pose[:3] = 0 item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0) if 'mesh_body_posed' in self.signature: mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) item['mesh_body_posed'] = mesh_body_posed_no_orient if 'mesh_posed' in self.signature: mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32) item['mesh_posed'][mesh_mask] = mesh_posed if 'mesh_unposed' in self.signature: mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32) item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape if 'mesh_normals' in self.signature: mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32) item['mesh_normals'][mesh_mask] = mesh_normals for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_normals' + res_string in self.signature: uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals'] item['uv_normals' + res_string] = uv_normals if 'uv_unposed' + res_string in self.signature: uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in self.signature: uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string[1:], f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] item['uv_body_posed' + res_string] = uv_body_posed_no_orient if 'uv_posed' + res_string in self.signature: uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed'] item['uv_posed' + res_string] = uv_posed data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()} for key, value in data.items(): if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): data[key] = torch.permute(value, [2,0,1]) if self.cache is not None: self.cache[frame] = data data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} data = { **data, **self.static_cache[frame] } return data def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): if not is_cape: return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False) else: return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True) if __name__ == '__main__': from torch.utils.data import DataLoader import plotly.graph_objects as go import plotly.express as ex from plotly.subplots import make_subplots from custom_utils.custom_utils import _maskFaces from losses.losses import collision_loss, normal_loss from smpl.smpl_torch import SMPL from torch.profiler import profile, record_function, ProfilerActivity F = np.load('smpl/smpl_faces.npy') F_skirt = np.load('smpl/smpl_faces_skirt.npy') device = "cuda" if torch.cuda.is_available() else "cpu" signature = { 'pose_shape', 'uv_posed', 'uv_body_posed', 'mesh_body_posed', 'mesh_posed', 'mesh_unposed', 'uv_unposed', 'uv_normals', 'mesh_normals', 'uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', } dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) smpl = SMPL(orig_smpl=False, path='smpl', device='cuda') # create a profiler object profiler = torch.profiler.profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=2, warmup=1, active=3, repeat=2), on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"), with_stack=True, profile_memory=True, ) step = 0 for batch in train_dataloader: batch print(batch['id'][0]) with profiler: m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0] profiler.step() step += 1 if step > 25: break profiler.stop() print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25)) print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=25)) print(profiler.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=10)) # b = 0 # mesh_mask = batch['mesh_mask'][b].cpu().numpy() # faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F # _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) # c = collision_loss(batch['mesh_posed'], batch['mesh_body_posed'], batch['mesh_mask']) # fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), # go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,1], z=batch['mesh_body_posed'].cpu().numpy()[0,:,2], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1), # ]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() # mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] # mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] # mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] # mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() # fig = make_subplots(rows=2, cols=2, # specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}], # [{'type': 'mesh3d'}, {'type': 'mesh3d'}]]) # fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1) # fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2) # fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1) # fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2) # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False), # yaxis=dict(range=[-1, 1], autorange=False), # zaxis=dict(range=[-1, 1], autorange=False), # aspectmode='cube'), # scene2=dict(xaxis=dict(range=[-1, 1], autorange=False), # yaxis=dict(range=[-1, 1], autorange=False), # zaxis=dict(range=[-1, 1], autorange=False), # aspectmode='cube'), # scene3=dict(xaxis=dict(range=[-1, 1], autorange=False), # yaxis=dict(range=[-1, 1], autorange=False), # zaxis=dict(range=[-1, 1], autorange=False), # aspectmode='cube'), # scene4=dict(xaxis=dict(range=[-1, 1], autorange=False), # yaxis=dict(range=[-1, 1], autorange=False), # zaxis=dict(range=[-1, 1], autorange=False), # aspectmode='cube')) # fig.show() # uv_mask = batch['uv_mask'][b].cpu().numpy() # uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) # uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 # uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 # uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2 # uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2 # uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3 # img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2) # img = np.transpose(img, (1, 2, 0)) # ex.imshow(img).show() # pass pass # if __name__ == '__main__': # os.environ["CUDA_VISIBLE_DEVICES"] = "6" # from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces # import plotly.graph_objects as go # import plotly.express as ex # # import cmapy as cm # uv2mesh = UV2Mesh(is_skirt=True) # F = np.load("custom_utils/smpl_faces.npy") # sum_uv = np.zeros((512, 256, 3), dtype=np.float32) # sum_mask = np.zeros((512, 256), dtype=np.float32) # for i, batch in enumerate(dataset): # print(i) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # ex.imshow(batch['uv_normals'][0].numpy().argmax(-1) / 6.0).show() # sum_uv += batch['uv_unposed'].numpy().sum(axis=0) # sum_mask += batch['uv_mask'].numpy().sum(axis=0) # sum_mask[sum_mask == 0] = 1 # avg_garment = sum_uv / sum_mask[:, :, np.newaxis] # # x = np.transpose(avg_garment, [1,2,0]) # x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0] # 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 = F[:,0], # j = F[:,1], # k = F[:,2] # )]).show() # np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment) # print()