|
- import logging
- import os
- import glob
- import numpy as np
- # import tensorflow.compat.v1 as tf
- import tensorlayer as tl
- import argparse
- import importlib
- from tqdm import trange
- from utils import grid
-
- logging.basicConfig(
- level=logging.INFO,
- format='%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s',
- datefmt="%Y-%m-%d %H:%M:%S")
- logger = logging.getLogger(__name__)
-
- RANDOM_SEED = 42
- np.random.seed(RANDOM_SEED)
- # tf.set_random_seed(RANDOM_SEED)
-
- def gen_folding():
- Trainer = getattr(importlib.import_module('11_train'), 'Trainer')
- trainer = Trainer(args.checkpoint_dir, args.input_pattern,args.grid_steps,args.model,1,train_flag=False)
- files = glob.glob(args.input_pattern, recursive=True)
- if len(files) > 1:
- origin_folder = os.path.commonpath(files)
- elif '*' in files[0]:
- origin_folder = files[0][:args.input_pattern.find('*') - 1]
- else:
- origin_folder = os.path.split(files[0])[0]
- assert len(files) > 0, "No files found"
-
- grid_steps = trainer.grid_steps #[218, 353, 1]跟点的个数n有关系
- model = trainer.model
-
- for i in trange(len(files)):
- logger.info(f'Run network for {files[i]}')
- # Run network
- x_tilde = model(trainer.x)
- col_loss = trainer.eval_loss(x_tilde) #颜色从原始点映射到colors_tilde,再映射回原始点,并转为yuv格式,计算原始yuv与映射回原始yuv的差方
- # Obtain data
- ori_values, norm_params = trainer.data[i]
-
- # Paths and folders
- origin_path = files[i]
- filepath, filename = os.path.split(origin_path[len(origin_folder) + 1:])
- current_output_dir = os.path.join(args.output_dir, filepath)
- os.makedirs(current_output_dir, exist_ok=True)
- output_file = os.path.join(current_output_dir, 'folding_data') #/userhome/pcc_attr_folding/pcc_attr_folding-master/manual_patches/redandblack_vox10_1550/eval/00/folding_data.npz
-
- np.savez_compressed(output_file,
- x_tilde=x_tilde, ori_values=ori_values, norm_params=norm_params,
- grid_steps=grid_steps, origin_path=origin_path, origin_folder=origin_folder)
-
- if __name__ == '__main__':
- parser = argparse.ArgumentParser(
- prog='20_gen_folding.py',
- description='Compute the folding using a trained model.',
- formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- parser.add_argument('input_pattern', help='Input pattern.')
- parser.add_argument('output_dir', help='Output directory.')
- parser.add_argument('checkpoint_dir', help='Directory where to save/load model checkpoints.')
- parser.add_argument('--model', help='Model module.', required=True)
- parser.add_argument('--input_pipeline', help='Input pipeline module.', required=True)
- parser.add_argument('--filters', type=int, default=32, help='Number of filters per layer.')
- parser.add_argument('--max_steps', type=int, default=100000, help='Train up to this number of steps.')
- parser.add_argument('--grid_steps', default='512,512,1')
- parser.add_argument('--grid_steps_factor', default=1.0, type=float, help='Multiply grid_steps by this factor.')
- args = parser.parse_args()
-
- os.makedirs(os.path.split(args.checkpoint_dir)[0], exist_ok=True)
- os.makedirs(os.path.split(args.output_dir)[0], exist_ok=True)
- gen_folding()
|