import os import sys # os.environ["CUDA_VISIBLE_DEVICES"] = '6' 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 torch.utils.data import Dataset 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 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.filled = '' 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): data = dict() 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, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; 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, 'uv_mask' + res_string, id + '.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', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32) if 'mesh_mask' in self.signature: if self.is_cape: data['mesh_mask'] = np.ones((3638), np.float32) else: mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask # 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, id) -> dict: if id not in self.store: item = self.read(id) 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[id] = stored_data data = { key: value.to(self.device) for key, value in self.store[id].items()} return data class Dataset(Dataset): def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', 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(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) self.body_faces = np.load('smpl/smpl_faces.npy') self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None 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.id] } return data separator = '_' if self.is_cape: separator = '.' uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' if self.is_cape: uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' item = dict() if 'id' in self.signature: item['id'] = frame.id if 'frame_nr' in self.signature: item['frame_nr'] = frame.frame_nr if 'pose_shape' in self.signature: pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32) pose_shape[:3] = 0 item['pose_shape'] = pose_shape if 'mesh_body_posed' in self.signature: mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] item['mesh_body_posed'] = mesh_posed_body_no_orient if 'mesh_posed' in self.signature: mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32) item['mesh_posed'] = mesh_posed if 'mesh_unposed' in self.signature: mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] item['mesh_unposed'] = mesh_unposed_no_shape if 'mesh_curvature' in self.signature: mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] item['mesh_curvature'] = mesh_curvature.astype(np.int32) if 'mesh_normals' in self.signature: mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals'] item['mesh_normals'] = mesh_normals for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_curvature' + res_string in self.signature: uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); if res_string != '': uv_curvature = uv_curvature['uv_curvature'] else: uv_curvature = uv_curvature['curvature_uv'] uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device) item['uv_curvature' + res_string] = uv_curvature 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, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in self.signature: uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz')) uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] item['uv_body_posed' + res_string] = uv_posed_body_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, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] item['uv_posed' + res_string] = uv_posed if 'uv_normals' + res_string in self.signature: uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals'] item['uv_normals' + res_string] = uv_normals data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) 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.id] } return data if __name__ == '__main__': from torch.utils.data import DataLoader device = "cuda" if torch.cuda.is_available() else "cpu" signature = { 'id', 'frame_nr', 'pose_shape', 'uv_posed', 'uv_body_posed', 'mesh_body_posed', 'mesh_posed', 'mesh_unposed', 'uv_unposed', 'uv_curvature', 'mesh_curvature', 'uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body', } dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: batch 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_curvature'][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()