From 1d6a089dd4eb60916de96faef6e65a8b4aed6688 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 1 Apr 2023 19:07:53 +0000 Subject: [PATCH] new dataloader --- data/dataset.py | 662 +++++++++++++++--------------------------------- data/utils.py | 6 +- 2 files changed, 207 insertions(+), 461 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index 1e941c5..c667724 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -1,79 +1,101 @@ import os import sys + +os.environ["CUDA_VISIBLE_DEVICES"] = '6' + +import math import time -import cv2 +import torch import random -import pickle +import pytorch3d +import torchvision as tv import numpy as np -import tensorflow as tf 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 custom_utils.custom_utils import nest_conditions, batch_to_dict -from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame +from data.utils import Frame -class StaticDataset(): - def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> None: +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] + else: + frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + + 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], in_memory=True, is_cape=False) -> None: self.store = dict() self.signature = signature self.folder = folder self.in_memory = in_memory assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) - self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256', - 'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256', - 'mesh_static', 'mesh_mask', 'uv_shaped_body'] - self.keys = [value for value in signature.keys() if value in self.possible_keys] - if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys: - self.keys.append('uv_static_128') self.resolutions = resolutions self.filled = '' self.is_cape = is_cape - self.use_no_orient = use_no_orient - + if is_cape: self.uv_masks = dict() uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32) self.uv_masks[512] = uv_mask.astype(np.float32) for res in [32, 64, 128, 256]: - down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) + down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST) down = np.squeeze(down.numpy().astype(np.float32)) self.uv_masks[res] = down - def read(self, id, garment): + def read(self, id): data = dict() - info = loadInfo(os.path.join(self.folder, id)) - - data['shape'] = info['shape'] - data['poses'] = info['poses'] - - data['map'] = sio.loadmat(os.path.join(self.folder, id, garment + '_map.mat'))['map'].reshape((-1)) - - - if 'mesh_mask' in self.signature.keys(): - if self.is_cape: - data['mesh_mask'] = np.ones((3638), np.float32) - else: - data['mesh_mask'] = read_garment_mask(self.folder, id, garment).astype(np.float32) - for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) - if 'uv_static' + res_string in self.signature.keys(): + 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.keys(): + 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.keys(): - data['mesh_static'] = read_garment_template(self.folder, id, garment) + 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 @@ -90,462 +112,186 @@ class StaticDataset(): def expand_batch(self, batch): - for key in self.keys: - batch[key] = tf.zeros(((0,) + self.signature[key])) - for i in range(batch['id'].shape[0]): - item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8')) - for key in self.keys: - if key == 'pose': - frame_nr = int(batch['frame_nr'][i].numpy().decode('UTF-8')) - pose = item[key][frame_nr, :] - if self.use_no_orient: - pose[:3] = 0 - batch[key] = tf.concat((batch[key], tf.expand_dims(pose, axis=0)), axis=0) - else: - batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0) + batch_size = batch['id'].shape[0] + items = [] + for i in range(batch_size): + items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8'))) + for key in items[0].keys(): + batch[key] = torch.stack([item[key] for item in items], dim=0) return batch +class Dataset(Dataset): + def __init__(self, folder, signature, resolutions, 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 -class DatasetIdx(tf.data.Dataset): + self.transforms = tv.transforms.ToTensor() + self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) - def __new__(cls, 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)) + self.body_faces = np.load('smpl/smpl_faces.npy') + self.static_cache = StaticCache(folder, signature, is_cape=is_cape) - # if is_cape: - # frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] - # else: - # frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] - - # if testing: - # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) - # frames = [frame.split('.')[0] + '_4' for frame in frames] - - frames = [] - with open(os.path.join(folder, frames_file), 'r') as f: - lines = f.readlines() - lines = [line.strip() for line in lines] - frames = lines - - if shuffle: - random.shuffle(frames) - - return tf.data.Dataset.from_tensor_slices(frames) + self.cache = dict() if use_cache else None -class DatasetFactory(): + def __len__(self): + return len(self.frames) - def __init__(self): - super(DatasetFactory, self).__init__() - def _build_tf_signature(self, signature): - tf_signature = [] - for key in signature.keys(): - if key == 'id': - tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) - elif key == 'frame_nr': - tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) - elif key == 'garment': - tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) - elif key == 'mesh_curvature': - tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32)) - else: - tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) - return tf_signature - - - def get_dataset(self, - folder, - batch_size, - signature, - shuffle=True, - resolutions=[512], - in_memory=False, - do_repeat=True, - is_cape=False, - use_no_orient=True, - testing=False, - frames_file=None): - assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) - - dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file) + def __getitem__(self, idx): + frame : Frame = self.frames[idx] + if self.cache is not None: + if frame in self.cache.keys(): + return self.cache[frame] separator = '_' - if is_cape: + if self.is_cape: separator = '.' - uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' + 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' - if use_no_orient == False: - assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!' + item = dict() - to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) - shapedirs = [] - with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f: - female = pickle.load(f, encoding="latin1") - shapedirs.append(female['shapedirs'][to_keep, :]) - with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f: - male = pickle.load(f, encoding="latin1") - shapedirs.append(male['shapedirs'][to_keep, :]) + 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')) + 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) + 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).to(self.device) + item['uv_curvature' + res_string] = uv_curvature - def read_item(frame): - frame_split = tf.strings.split(frame, sep=separator) - id, frame_nr, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8") - item = dict() + 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 - # redunda ntly loading info again - pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr] - # !!!!!!!!!!!!!!!!!!! - map = None + 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 'id' in signature.keys(): - item['id'] = id - if 'frame_nr' in signature.keys(): - item['frame_nr'] = frame_nr - if 'garment' in signature.keys(): - item['garment'] = garment + 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 'mesh_body_posed' in signature.keys(): - if use_no_orient: - item['mesh_body_posed'] = read_posed_body_frame(folder, id, frame_nr, pose, deorient=use_no_orient) - if 'mesh_posed' in signature.keys(): - if use_no_orient: - item['mesh_posed'] = read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False) - # if 'mesh_unposed' in signature.keys(): - # mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] - # item['mesh_unposed'] = mesh_unposed_no_shape + data = { **item, **self.static_cache[frame.id] } - if 'mesh_curvature' in signature.keys(): - mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] - item['mesh_curvature'] = mesh_curvature.astype(np.int32) - - - for res in resolutions: - res_string = '' - if res != 512: - res_string = '_' + str(res) - - if 'uv_curvature' + res_string in signature.keys(): - uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz')); - if res_string != '': - uv_curvature = uv_curvature['uv_curvature'] - else: - uv_curvature = uv_curvature['curvature_uv'] - uv_curvature = tf.one_hot(uv_curvature, 7) - item['uv_curvature' + res_string] = uv_curvature - - if 'uv_unposed' + res_string in signature.keys(): - uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_string, id + separator + 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 signature.keys(): - if use_no_orient: - uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + 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 - else: - uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body'] - item['uv_body_posed' + res_string] = uv_body_posed - - if 'uv_posed' + res_string in signature.keys(): - if use_no_orient: - uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] - item['uv_posed' + res_string] = uv_posed - else: - uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt'] - item['uv_posed' + res_string] = uv_posed - - as_list = [] - for key in signature.keys(): - if key in item: - as_list.append(item[key]) - - return as_list - - - if in_memory: - dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, - inp=[i], - Tout=self._build_tf_signature(signature) - ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True) - else: - dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, - inp=[i], - Tout=self._build_tf_signature(signature) - ), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True) - - if do_repeat: - dataset = dataset.repeat() - - return dataset + data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()} + if self.cache is not None: + self.cache[frame] = data + return data 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") - - - # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train') - # for d in dataset: - # print(d) - - factory = DatasetFactory() - signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (), - # 'pose_shape': (72 + 11,), - # 'uv_posed': (3, 512, 256,), - # 'uv_body_posed': (512, 256, 3), - # 'uv_body_posed_32': (32, 16, 3), - # 'mesh_body_posed': (14475, 3,), - # 'mesh_posed': (14475, 3,), - # 'mesh_unposed': (14475, 3,), - 'uv_unposed': (512, 256, 3), - # 'uv_unposed_32': (3, 32, 16,), - 'uv_curvature': (512, 256, 7), - # 'mesh_curvature': (14475) - } - signature_static = { #'uv_static': (3, 512, 256,), - # 'uv_static_32': (3, 32, 16,), - # 'mesh_static': (14475, 3,), - # 'mesh_mask': (14475,), - 'uv_mask': (512, 256,), - # 'uv_shaped_body': (3, 512, 256,), - } - - dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt', - batch_size=1, - signature=signature_dynamic, - resolutions=[512], - do_repeat=False, - use_one_hot=False, - is_cape=False, - frames_file='frames_dress_30k.txt') - static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt', - signature=signature_static, - resolutions=[512], - is_cape=False) - - - # for i, batch in enumerate(dataset): - # print(i) - - - 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() - - # pass - - - # uv_res = 256 - - # 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 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 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 - - - # uv_map = np.load('custom_utils/stretched_uv_map_3.npy') - # m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() - # m2uv = m2uv_array(m2uv) - # uv_mask = np.load('custom_utils/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 = 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] - - - - # for i, batch in enumerate(dataset): - # print(i) - # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) - - # uv_unposed = np.squeeze(batch['uv_unposed'].numpy()) - # mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) - # mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) - - - # values = mesh_unposed - # uv_values = np.zeros((uv_map.shape[0], 3)) - # uv_values[m2uv[:, 0]] = values - # uv_values[m2uv[:, 1]] += values - # uv_counter = np.zeros((uv_map.shape[0])) - # uv_counter[m2uv[:, 0]] = 1 - # uv_counter[m2uv[:, 1]] += 1 - # uv_values /= uv_counter[:, np.newaxis] - # x = uv_values - # # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) - # im = ((x - x.min()) / (x.max() - x.min())) * 255 - - - # x = np.transpose(uv_unposed, [1,2,0]) - # x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :] - # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) - # 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 = faces[:,0], - # j = faces[:,1], - # k = faces[:,2] - # )]).show() - - - # print() - - - - - # # @tf.function - # def euclid_error(mask, gt, pred): - # # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1) - # return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)) - - # all_error = [] - # for batch in dataset: - # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) - # mesh_reconstruct = uv2mesh(batch['uv_posed']) - # err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed'])) - - # i = 296 - # err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min()) - # colors = np.array([cm.color('jet', 1 - i) for i in err_normed]) - - # _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool)) - # x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :] - # y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :] - # colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :] - # fig = 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 = faces[:,0], - # j = faces[:,1], - # k = faces[:,2] - # ), - # go.Mesh3d( - # # 8 vertices of a cube - # x=y[:,0], - # y=y[:,1], - # z=y[:,2], - # # i, j and k give the vertices of triangles - # i = faces[:,0], - # j = faces[:,1], - # k = faces[:,2], - # )]) - # fig.show() - - # fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2]) - # fig.update_traces(marker={'color': colors, 'size': 5}) - # fig.show() - - # all_error.append(err.numpy().sum()) - # # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic))) - # print(len(all_error)) - + from torch.utils.data import DataLoader - # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2) + device = "cuda" if torch.cuda.is_available() else "cpu" - # from losses import collision_loss as collision_penalty - # from losses import collision_penalty as collision_penalty_2 + 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 - # for d in dataset: - # d - # v = uv2mesh(d[0]) - # b = d[3] - # m = d[4] - # b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0) - # c = collision_penalty(v, b, m) - # c2 = collision_penalty(v, b2, m) - # c_ = collision_penalty_2(v, b, m) - # c2_ = collision_penalty_2(v, b2, m) - # print() - # pass - \ No newline at end of file + + + +# 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() \ No newline at end of file diff --git a/data/utils.py b/data/utils.py index e7d6392..6f25887 100644 --- a/data/utils.py +++ b/data/utils.py @@ -1,16 +1,16 @@ import tables as tb class Frame: - def __init__(self, id, frame_nr, garment): + def __init__(self, id, frame_nr, garment=''): self.id = id self.frame_nr = frame_nr self.garment = garment def __hash__(self): - return hash((self.id, self.frame_nr)) + return hash((self.id, self.frame_nr, self.garment)) def __eq__(self, other): - return (self.id, self.frame_nr) == (other.id, other.frame_nr) + return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment) def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):