diff --git a/configs/iccvw/abl_all.py b/configs/iccvw/abl_all.py new file mode 100644 index 0000000..b9d9639 --- /dev/null +++ b/configs/iccvw/abl_all.py @@ -0,0 +1,195 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'abl_iccvw_all_1' +config['notes'] = '' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = '' + +config['use_skinning'] = True # and deshaping +config['use_normals'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh + +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['avg_garment_path'] != '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['static_model'] = 'checkpoints/iccvw/static_iccvw_4' +config['body_model'] = 'checkpoints/iccvw/body_iccvw_1' +config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 +config['balancing_activation_limit'] = 0.01 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_normals', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_norm ############################ +config['cond_ae_norm'] = dict() +config['cond_ae_norm']['device'] = config['device'] +config['cond_ae_norm']['latent_dim'] = config['latent_dim'] +config['cond_ae_norm']['kernel_size'] = 7 +config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_norm']['encoder_block_settings'] = block_setting +config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_norm']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/configs/iccvw/ae_norm_1.py b/configs/iccvw/ae_norm_1.py index f4cdb16..c9c1383 100644 --- a/configs/iccvw/ae_norm_1.py +++ b/configs/iccvw/ae_norm_1.py @@ -58,7 +58,6 @@ block_setting_reverse = [ CNBlockConfig(64, 32, 9), CNBlockConfig(32, 16, 3), CNBlockConfig(16, 3, 3), - # CNBlockConfig(3, 3, 3), ] config['cond_ae']['encoder_block_settings'] = block_setting diff --git a/custom_data/avg/train_no_shape_avg_uv.npy b/custom_data/avg/train_no_shape_avg_uv.npy index 7868198..f2c9cd3 100644 Binary files a/custom_data/avg/train_no_shape_avg_uv.npy and b/custom_data/avg/train_no_shape_avg_uv.npy differ diff --git a/custom_data/avg/train_no_shape_skirt_avg_uv.npy b/custom_data/avg/train_no_shape_skirt_avg_uv.npy deleted file mode 100644 index ac60c6e..0000000 Binary files a/custom_data/avg/train_no_shape_skirt_avg_uv.npy and /dev/null differ diff --git a/custom_data/avg/train_skirt_no_shape_avg_uv.npy b/custom_data/avg/train_skirt_no_shape_avg_uv.npy new file mode 100644 index 0000000..ba7d72d Binary files /dev/null and b/custom_data/avg/train_skirt_no_shape_avg_uv.npy differ diff --git a/data/preprocessing.py b/data/preprocessing.py index 9a4a606..dff4038 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -679,7 +679,8 @@ def check_irregular_poses(folder, folder_reg, out_folder): # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') - mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) + mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] + mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr] F = np.load('smpl/smpl_faces_skirt.npy') _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) @@ -694,6 +695,46 @@ def check_irregular_poses(folder, folder_reg, out_folder): fig.show() pass + +def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'): + from custom_utils.custom_utils import UV2Mesh + import plotly.express as ex + import plotly.graph_objects as go + + frames = [] + with open(os.path.join(out_folder, frames_file), 'r') as f: + for line in f: + line = line.strip() + id, frame_nr, garment = line.split(' ') + frames.append(Frame(id, int(frame_nr), garment)) + + uv_sum = np.zeros((512, 256, 3), dtype=np.float32) + mask_sum = np.zeros((512, 256), dtype=np.float32) + + for frame in tqdm(frames): + uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed'] + mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask'] + uv_sum += uv + mask_sum += mask + + mask_sum[mask_sum == 0] = 1 + avg_garment = uv_sum / mask_sum[:, :, None] + + avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0] + F = np.load('smpl/smpl_faces.npy') + + fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,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() + + np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment) + + + if __name__ == '__main__': @@ -721,9 +762,10 @@ if __name__ == '__main__': # 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) + calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) - create_frames_file_partial(folder, out_folder, part=0.5) + # create_frames_file_partial(folder, out_folder, part=0.5) # folder = '/data-local/127_2/cape_processed/test/female' # create_mesh_normals_cape(folder, None) \ No newline at end of file diff --git a/models/cvae_model.py b/models/cvae_model.py index b81ee6a..788f31e 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -18,14 +18,13 @@ class CVAE(nn.Module): super(CVAE, self).__init__() self.size = config['resolution'] - if not('only_curvatures' in config.keys() and config['only_curvatures'] == True): - self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device']) - if config['static_model'] != '': - load_model(self.static_enc, config['static_model'] + '_enc') + self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device']) + if config['static_model'] != '': + load_model(self.static_enc, config['static_model'] + '_enc') - self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device']) - if config['body_model'] != '': - load_model(self.body_enc, config['body_model'] + '_enc') + self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device']) + if config['body_model'] != '': + load_model(self.body_enc, config['body_model'] + '_enc') if config['use_normals']: self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device']) diff --git a/scripts/ae_norm_train.py b/scripts/ae_norm_train.py index 597e662..af79d1c 100644 --- a/scripts/ae_norm_train.py +++ b/scripts/ae_norm_train.py @@ -205,15 +205,15 @@ if __name__ == '__main__': load_model(model.encoder, config['pretrained_model'] + '_enc') load_model(model.decoder, config['pretrained_model'] + '_dec') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5) + ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5) + ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5) run_name = config['name'] print(f'started training: {run_name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=10000, log_graph=True) + wandb.watch(model, log='all', log_freq=20000, log_graph=True) train() diff --git a/scripts/cvae.py b/scripts/cvae.py index c485b5e..2ee71e9 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -10,10 +10,12 @@ parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) -config_module = importlib.import_module(args['config']) -config = config_module.config + if args['debug']: - from configs.abl_norm import config + from configs.iccvw.abl_all import config +else: + config_module = importlib.import_module(args['config']) + config = config_module.config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] @@ -42,6 +44,7 @@ from data.dataset import Dataset from models.cond_encoder import AutoEncoder from utils.checkpoints import CheckpointSaver, load_model from utils.utils import InfiniteDataLoader as DataLoader +from utils.utils import text_to_html from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss from pytorch_msssim import MS_SSIM @@ -51,7 +54,8 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary def init_wandb(): - run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/cvae.py') artifact.add_file('losses/losses.py') @@ -255,7 +259,14 @@ def plot_everything(batch: dict, type: str): mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0)) logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0)) - return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } + batch_id = batch['id'] + batch_frame_nr = batch['frame_nr'].cpu().numpy() + batch_garment = batch['garment'] + batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))] + batch_info = text_to_html(batch_info) + batch_info = wandb.Html(batch_info) + + return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } def sample(folder): @@ -455,10 +466,10 @@ def train(): valid_accumulator.reset() print(step) - if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017): + if train_loss['train/euclid_error'] < config['balancing_activation_limit']: g_balancing_coeff += 1 / 10000 - if step % 10000 == 0: + if step % 25000 == 0: ckpt_saver(model, step, optimizer, None) step += 1 @@ -485,7 +496,7 @@ if __name__ == '__main__': _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) g_balancing_coeff = config['starting_balancing_coeff'] - ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=10) step = 0 if config['pretrained_model'] != '': diff --git a/utils/checkpoints.py b/utils/checkpoints.py index 30db495..b564e40 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -39,5 +39,6 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No model.load_state_dict(checkpoint['model']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict()) if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler']) + print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}') return int(iteration) if iteration != '' else 0 diff --git a/utils/utils.py b/utils/utils.py index d1d247a..5b5fe02 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -42,4 +42,10 @@ class InfiniteDataLoader(DataLoader): # Dataset exhausted, use a new fresh iterator. self.dataset_iterator = super().__iter__() batch = next(self.dataset_iterator) - return batch \ No newline at end of file + return batch + +def text_to_html(lines): + html = '' + for line in lines: + html += "
" + line + "
" + return html \ No newline at end of file