|
- #!/usr/bin/env python
- # coding: utf-8
-
- # Copyright (c) Nanjing University, Vision Lab.
- # Last update:
- # 2019.10.27
- # 2019.11.14
- # 2020.11.26
-
- import os
- import time
- import numpy as np
- import tensorflow as tf
- import matplotlib.pylab as plt
- import pandas as pd
- import subprocess
- import glob
- import configparser
- import argparse
- import importlib
- # from numba import cuda
- tf.enable_eager_execution()
-
- from process import preprocess, postprocess
- # import models.model_voxception as model
- from transform import compress_factorized, decompress_factorized
- from transform import compress_hyper, decompress_hyper
-
- from dataprocess.inout_bitstream import write_binary_files_factorized, read_binary_files_factorized
- from dataprocess.inout_bitstream import write_binary_files_hyper, read_binary_files_hyper
-
- os.environ['CUDA_VISIBLE_DEVICES']="0"
- # set gpu.
- cfg = tf.ConfigProto()
- cfg.gpu_options.per_process_gpu_memory_fraction = 1.0
- cfg.gpu_options.allow_growth = True
- cfg.log_device_placement=True
- # config.device_count={'gpu':0}
- sess = tf.Session(config=cfg)
-
- from myutils.pc_error_wrapper import pc_error
- from myutils.pc_error_wrapper import get_points_number
-
-
- def set_default_config(input_file, cfg_rootdir, resolution, mode, cube_size, modelname="models.model_voxception"):
- filename = os.path.split(input_file)[-1][:-4]
- config_file = os.path.join(cfg_rootdir, filename+'.ini')
- config = configparser.ConfigParser()
- if os.path.exists(config_file):
- config.read(config_file)
- print('config already exists.')
- else:
- config["DEFAULT"] = {"cube_size": cube_size, "min_num": 64, "resolution":resolution}
- if mode=="factorized":
- if modelname=="models.model_voxception":
- config["R1"] = {"scale": 0.625, "ckpt_dir": './checkpoints/factorized/a2b3/'}
- config["R2"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/a2b3/'}
- config["R3"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/a4b3/'}
- config["R4"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/a6b3/'}
- config["R5"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/a10b3/'}
- config["R6"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/a16b3/'}
- if modelname=="models.model_simple":
- config["R1"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a1b3/'}
- config["R2"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a2b3/'}
- config["R3"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a3b3/'}
- config["R4"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a4b3/'}
- config["R5"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a5b3/'}
- config["R6"] = {"scale": 1.0, "ckpt_dir": './checkpoints/factorized/simple/a6b3/'}
-
- elif mode=="hyper":
- config["R1"] = {"scale": 5/8., "ckpt_dir": './checkpoints/hyper/a0.75b3/'}
- config["R2"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a0.75b3/'}
- config["R3"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a2b3/'}
- config["R4"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a3.5b3/'}
- config["R5"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a6b3/'}
- config["R6"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a10b3/'}
- config["R7"] = {"scale": 1.0, "ckpt_dir": './checkpoints/hyper/a16b3/'}
- config.write(open(config_file, 'w'))
- print('initialize config.')
- return config, config_file
-
-
- def test_factorized(input_file, model, ckpt_dir, scale, cube_size, min_num, postfix=''):
- # Pre-process
- cubes, cube_positions, points_numbers = preprocess(input_file, scale, cube_size, min_num)
- ### Encoding
- strings, min_v, max_v, shape = compress_factorized(cubes, model, ckpt_dir)
- # Write files
- filename = os.path.split(input_file)[-1][:-4]
- print(filename)
- rootdir = './compressed'+ postfix +'/'
- bytes_strings, bytes_pointnums, bytes_cubepos = write_binary_files_factorized(
- filename, strings.numpy(), points_numbers, cube_positions,
- min_v.numpy(), max_v.numpy(), shape.numpy(), rootdir)
- # Read files
- strings_d, points_numbers_d, cube_positions_d, min_v_d, max_v_d, shape_d = \
- read_binary_files_factorized(filename, rootdir)
- # Decoding
- cubes_d = decompress_factorized(strings_d, min_v_d, max_v_d, shape_d, model, ckpt_dir)
-
- # bpp
- N = get_points_number(input_file)
- bpp = round(8*(bytes_strings + bytes_pointnums + bytes_cubepos)/float(N), 4)
- bpp_strings = round(8*bytes_strings/float(N), 4)
- bpp_pointsnums = round(8*bytes_pointnums/float(N) ,4)
- bpp_cubepos = round(8*bytes_cubepos/float(N), 4)
- bpp_strings_hyper = 0
- bpp_strings_head = 0
- bpps = [bpp, bpp_strings, bpp_strings_hyper, bpp_strings_head, bpp_pointsnums, bpp_cubepos]
-
- return cubes_d, cube_positions_d, points_numbers_d, N, bpps
-
-
- def test_hyper(input_file, model, ckpt_dir, scale, cube_size, min_num, postfix=''):
- # Pre-process
- cubes, cube_positions, points_numbers = preprocess(input_file, scale, cube_size, min_num)
- ### Encoding
- y_strings, y_min_vs, y_max_vs, y_shape, z_strings, z_min_v, z_max_v, z_shape, x_ds = compress_hyper(cubes, model, ckpt_dir, True)
- # Write files
- filename = os.path.split(input_file)[-1][:-4]
- print(filename)
- rootdir = './compressed'+ postfix +'/'
- bytes_strings, bytes_strings_head, bytes_strings_hyper, bytes_pointnums, bytes_cubepos = write_binary_files_hyper(
- filename, y_strings.numpy(), z_strings.numpy(), points_numbers, cube_positions,
- y_min_vs.numpy(), y_max_vs.numpy(), y_shape.numpy(),
- z_min_v.numpy(), z_max_v.numpy(), z_shape.numpy(), rootdir)
- # Read files
- y_strings_d, z_strings_d, points_numbers_d, cube_positions_d, y_min_vs_d, y_max_vs_d, y_shape_d, z_min_v_d, z_max_v_d, z_shape_d = \
- read_binary_files_hyper(filename, rootdir)
- # Decoding
- cubes_d = decompress_hyper(y_strings_d, y_min_vs_d.astype('int32'), y_max_vs_d.astype('int32'),
- y_shape_d, z_strings_d, z_min_v_d, z_max_v_d, z_shape_d, model, ckpt_dir)
- # cheat!!!
- ##############
- cubes_d = x_ds
- ##############
- # bpp
- N = get_points_number(input_file)
- bpp = round(8*(bytes_strings + bytes_strings_head + bytes_strings_hyper +
- bytes_pointnums + bytes_cubepos)/float(N), 4)
-
- bpp_strings = round(8*bytes_strings/float(N), 4)
- bpp_strings_hyper = round(8*bytes_strings_hyper/float(N), 4)
- bpp_strings_head = round(8*bytes_strings_head/float(N), 4)
- bpp_pointsnums = round(8*bytes_pointnums/float(N) ,4)
- bpp_cubepos = round(8*bytes_cubepos/float(N), 4)
- bpps = [bpp, bpp_strings, bpp_strings_hyper, bpp_strings_head, bpp_pointsnums, bpp_cubepos]
-
- return cubes_d, cube_positions_d, points_numbers_d, N, bpps
-
-
- def select_optimal_rho(item, rhos,
- input_file, output_file, input_file_n,
- cubes_d, points_numbers_d, cube_positions_d,
- scale, cube_size, res):
- for i, rho in enumerate(rhos):
- print('===== select rho =====')
- postprocess(output_file, cubes_d, points_numbers_d, cube_positions_d, scale, cube_size, rho)
- results = pc_error(input_file, output_file, input_file_n, res, show=False)
- PSNR = float(results[item])
- print('===== results: ', i, rho, item, PSNR)
- if i == 0:
- MAX_PSNR = 0
- optimal_rho = rho
- else:
- MAX_PSNR = max(PSNR, MAX_PSNR)
-
- if PSNR < MAX_PSNR:
- break
- else:
- optimal_rho = rho
-
- return optimal_rho
-
-
- def cfg_post_process(config, config_file, rate,
- input_file, output_file, input_file_n,
- cubes_d, points_numbers_d, cube_positions_d,
- scale, cube_size, res):
- # select rho for optimal d1/d2 metrics.
- if config.has_option(rate, 'rho_d1'):
- rho_d1 = float(config.get(rate, 'rho_d1'))
- else:
- item = "mseF,PSNR (p2point)"
- rhos = [0.8, 0.9, 1.0, 1.02, 1.05, 1.10, 1.15, 1.2, 1.25, 1.30, 1.40, 1.50, 1.75, 2.0, 2.5, 3.0]
- rho_d1 = select_optimal_rho(item, rhos,
- input_file, output_file, input_file_n,
- cubes_d, points_numbers_d, cube_positions_d,
- scale, cube_size, res)
- config.set(rate, 'rho_d1', str(rho_d1))
- config.write(open(config_file, 'w'))
-
- if config.has_option(rate, 'rho_d2'):
- rho_d2 = float(config.get(rate, 'rho_d2'))
- else:
- item = "mseF,PSNR (p2plane)"
- rhos = [1.0, 0.98, 0.95, 0.92, 0.90, 0.88, 0.85, 0.82, 0.80, 0.75, 0.70, 0.65, 0.50, 0.40, 0.30]
- rho_d2 = select_optimal_rho(item, rhos,
- input_file, output_file, input_file_n,
- cubes_d, points_numbers_d, cube_positions_d,
- scale, cube_size, res)
- config.set(rate, 'rho_d2', str(rho_d2))
- config.write(open(config_file, 'w'))
-
- return rho_d1, rho_d2
-
-
- def collect_results(results, results_d1, results_d2, bpps, N, scale, rho_d1, rho_d2):
- # bpp
- results["ori_points"] = N
- results["scale"] = scale
- # results["cube_size"] = cube_size
- # results["res"] = res
- results["bpp"] = bpps[0]
- results["bpp_strings"] = bpps[1]
- results["bpp_strings_hyper"] = bpps[2]
- results["bpp_strings_head"] = bpps[3]
- results["bpp_pointsnums"] = bpps[4]
- results["bpp_cubepos"] = bpps[5]
-
- results["rho_d1"] = rho_d1
- results["optimal D1 PSNR"] = results_d1["mseF,PSNR (p2point)"]
-
- results["rho_d2"] = rho_d2
- results["optimal D2 PSNR"] = results_d2["mseF,PSNR (p2plane)"]
-
- print(results)
-
- return results
-
-
- def plot_results(all_results, filename, root_dir):
- fig, ax = plt.subplots(figsize=(7.3, 4.2))
- plt.plot(np.array(all_results["bpp"][:]), np.array(all_results["mseF,PSNR (p2point)"][:]),
- label="D1", marker='x', color='red')
- plt.plot(np.array(all_results["bpp"][:]), np.array(all_results["mseF,PSNR (p2plane)"][:]),
- label="D2", marker='x', color = 'blue')
-
- plt.plot(np.array(all_results["bpp"][:]), np.array(all_results["optimal D1 PSNR"][:]),
- label="D1 (optimal)", marker='h', color='red', linestyle='-.')
- plt.plot(np.array(all_results["bpp"][:]), np.array(all_results["optimal D2 PSNR"][:]),
- label="D2 (optimal)", marker='h', color='blue', linestyle='-.')
- plt.title(filename)
- plt.xlabel('bpp')
- plt.ylabel('PSNR')
- plt.grid(ls='-.')
- plt.legend(loc='lower right')
- fig.savefig(os.path.join(root_dir, filename+'.png'))
-
- return
-
-
- def eval(input_file, rootdir, resolution, mode, cube_size, modelname, fixed_thres, postfix):
- # model = 'model_voxception'
- model = importlib.import_module(modelname)
-
- filename = os.path.split(input_file)[-1][:-4]
- output_file = filename + '_rec_' + postfix + '.ply'
- input_file_n = input_file
- csv_rootdir = os.path.join(rootdir, "csv")
- if not os.path.exists(csv_rootdir):
- os.makedirs(csv_rootdir)
- cfg_rootdir = os.path.join(rootdir, "cfg")
- if not os.path.exists(cfg_rootdir):
- os.makedirs(cfg_rootdir)
- # default config
- config, config_file = set_default_config(input_file, cfg_rootdir, resolution, mode, cube_size, modelname)
- cube_size = config.getint('DEFAULT', 'cube_size')
- min_num = config.getint('DEFAULT', 'min_num')
- res = config.getint('DEFAULT', 'resolution')
- print('cube size:', cube_size, 'min num:', min_num, 'res:', res)
-
- for index, rate in enumerate(config.sections()):
- scale = float(config.get(rate, 'scale'))
- ckpt_dir = str(config.get(rate, 'ckpt_dir'))
- print('====================', 'config:', rate, 'scale:', scale, 'ckpt_dir:', ckpt_dir)
-
- if mode=="factorized":
- cubes_d, cube_positions, points_numbers, N, bpps = test_factorized(input_file, model, ckpt_dir, scale, cube_size, min_num, postfix)
- elif mode == "hyper":
- cubes_d, cube_positions, points_numbers, N, bpps = test_hyper(input_file, model, ckpt_dir, scale, cube_size, min_num, postfix)
- cubes_d = cubes_d.numpy()
- print("bpp:",bpps[0])
- # select rho for optimal d1/d2 metrics.
- if fixed_thres==None:
- rho_d1, rho_d2 = cfg_post_process(config, config_file, rate,
- input_file, output_file, input_file_n,
- cubes_d, points_numbers, cube_positions,
- scale, cube_size, res)
- else:
- rho_d1, rho_d2 = 1.0, 1.0
-
- # metrics.
- rho = 1.0
- postprocess(output_file, cubes_d, points_numbers, cube_positions, scale, cube_size, rho, fixed_thres)
- results = pc_error(input_file, output_file, input_file_n, res, show=False)
-
- rho = rho_d1
- postprocess(output_file, cubes_d, points_numbers, cube_positions, scale, cube_size, rho, fixed_thres)
- results_d1 = pc_error(input_file, output_file, input_file_n, res, show=False)
-
- rho = rho_d2
- postprocess(output_file, cubes_d, points_numbers, cube_positions, scale, cube_size, rho, fixed_thres)
- results_d2 = pc_error(input_file, output_file, input_file_n, res, show=False)
-
- results = collect_results(results, results_d1, results_d2, bpps, N, scale, rho_d1, rho_d2)
-
- if index == 0:
- all_results = results.copy(deep=True)
- else:
- all_results = all_results.append(results, ignore_index=True)
-
- # write to csv
- print(all_results)
- if not os.path.exists(csv_rootdir):
- os.makedirs(csv_rootdir)
- csv_name = os.path.join(csv_rootdir, filename + '.csv')
- all_results.to_csv(csv_name, index=False)
- # plot
- plot_results(all_results, filename, csv_rootdir)
- return all_results
-
- def parse_args():
- parser = argparse.ArgumentParser(
- formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- parser.add_argument("--input", type=str, nargs='+', default='', dest="input")
- parser.add_argument("--rootdir", type=str, default='./results/', dest="rootdir")
- parser.add_argument("--resolution", type=int, default=1024, dest="resolution")
- parser.add_argument("--mode", type=str, default='hyper', dest="mode")
-
- # for ablation studies
- parser.add_argument("--cube_size", type=int, default=64, dest="cube_size")
- parser.add_argument("--modelname", default="models.model_voxception", help="(model_simple, model_voxception)", dest="modelname")
- parser.add_argument("--fixed_thres", type=float, default=None, help="fixed threshold ", dest="fixed_thres")
- parser.add_argument("--postfix", default="", help="", dest="postfix")
-
- args = parser.parse_args()
- print(args)
- return args
-
- if __name__ == "__main__":
- args = parse_args()
- if not os.path.exists(args.rootdir):
- os.makedirs(args.rootdir)
-
- # shapenet_filedirs = glob.glob("testdata/ShapeNet/*.ply")
- # args.input = shapenet_filedirs
- # print(shapenet_filedirs)
- for input_file in sorted(args.input):
- print(input_file)
- all_results = eval(input_file, args.rootdir, args.resolution, args.mode, args.cube_size, args.modelname, args.fixed_thres, args.postfix)
-
- """
- python eval_ablation_studies.py --input 'testdata/8iVFB/longdress_vox10_1300.ply' \
- 'testdata/8iVFB/loot_vox10_1200.ply' \
- 'testdata/8iVFB/redandblack_vox10_1550.ply' \
- 'testdata/8iVFB/soldier_vox10_0690.ply' \
- 'testdata/8iVFB/queen_0200.ply' \
- --rootdir='results/hyper/' \
- --mode="hyper" \
- --postfix='hyper'
-
- ablation studies:
- 1. cube_size:
- python eval_ablation_studies.py --input 'testdata/8iVFB/longdress_vox10_1300.ply'
- --rootdir='results/hyper/cube32'
- --mode="hyper"
- --postfix='hyper32'
- --cube_size=32
-
- python eval_ablation_studies.py --input 'testdata/8iVFB/longdress_vox10_1300.ply'
- --rootdir='results/hyper/cube128'
- --mode="hyper"
- --postfix='hyper128'
- --cube_size=128
-
- 2. compression netowork:
- python eval_ablation_studies.py --input 'testdata/8iVFB/longdress_vox10_1300.ply' \
- --rootdir='results/factorized/' \
- --mode="factorized" \
- --postfix='factorized'
-
- python eval_ablation_studies.py --input 'testdata/8iVFB/longdress_vox10_1300.ply' \
- --rootdir='results/factorized/simple_model/' \
- --mode="factorized" \
- --postfix='factorized_simple_model' \
- --model=models.model_simple
-
- 3. fixed threshold:
- python eveval_ablation_studiesal.py --input 'testdata/8iVFB/longdress_vox10_1300.ply' \
- --rootdir='results/hyper/fixed_thres/' \
- --mode="hyper" \
- --postfix='hyper_fixed' \
- --fixed_thres=0
- """
|