|
- # inputs: path to saved model, path to point clouds;
- # output: bit per occupied voxel
-
- import os
- os.environ['TL_BACKEND'] = 'tensorflow'
- import contextlib
- from arithmetic_coder import arithmetic_coding
- import numpy as np
- import argparse
- import time
- from utils.inout import occupancy_map_explore_test
- from voxel_dnn_coder.voxel_dnn_meta_endec import save_compressed_file,load_compressed_file
- import gzip
- import pickle
- from training.voxel_dnn_training_tensorlayer_2 import VoxelCNN
- #import tensorflow as tf
- import multiprocessing
- import tensorlayer as tl
- from tensorlayer.backend.ops.load_backend import BACKEND
-
- # encoding from breadth first sequence for parallel computing
- MAX_PROCESSES = 1 #原来4是最优的,但进程间会有冲突,加了锁也不行,只能设置为2
-
- def voxelCNN_decoding(args):
- pc_level, ply_path, output_path, model_path64, model_path32, model_path16, model_path8, bl_par_depth, signaling = args
- sequence_name = os.path.split(ply_path)[1]
- sequence = os.path.splitext(sequence_name)[0]
-
- output_path = output_path + str(sequence) + '/' + signaling
- os.makedirs(output_path, exist_ok=True)
- output_path = output_path + '/' + str(bl_par_depth) + 'levels'
- #outputfile = output_path + '.blocks.bin' #编码结果文件
- metadata_file = output_path + '.metadata.bin'
- heatmap_file = output_path + '.heatmap.pkl'
- start=time.time()
- #reading metadata
- with gzip.open(metadata_file, "rb") as f: #内容与block编码细节无关
- decoded_binstr,pc_level, departition_level =load_compressed_file(f)
- print('First decoded depth infor: ',pc_level,departition_level)
- #getting encoding input data
- boxes,binstr,no_oc_voxels,coor_min_max,lower_level_ocv=occupancy_map_explore_test(ply_path,pc_level,departition_level)
-
- #restore voxelCNN
- #voxelCNN64 = VoxelCNN(depth=64, height=64, width=64, residual_blocks=2, n_filters=64)
- #voxel_CNN64 = voxelCNN64.restore_voxelCNN(model_path64)
-
- '''voxelCNN32 = VoxelCNN(depth=32, height=32, width=32, residual_blocks=2, n_filters=64)
- voxel_CNN32 = voxelCNN32.restore_voxelCNN(model_path32)
-
- voxelCNN16 = VoxelCNN(depth=16, height=16, width=16, residual_blocks=2, n_filters=64)
- voxel_CNN16 = voxelCNN16.restore_voxelCNN(model_path16)
-
- voxelCNN8 = VoxelCNN(depth=8, height=8, width=8, residual_blocks=2, n_filters=32)
- voxel_CNN8 = voxelCNN8.restore_voxelCNN(model_path8)
- '''
- #voxel_CNN = [voxel_CNN64, voxel_CNN32, voxel_CNN16, voxel_CNN8]
- #voxel_CNN = [voxel_CNN64]
- #print(heatmap_file)
- with open(heatmap_file,'rb') as f: #内容是字典,key是block序号,需要与block对应上
- heatmap=pickle.load(f)
- #with open(outputfile, "rb") as inp:
- # bitin = arithmetic_coding.BitInputStream(inp)
- decoded_boxes=decompress_from_adaptive_freqs(boxes,heatmap, model_path64, output_path)
- end=time.time()
- decoding_time = end-start
- decoded_boxes=decoded_boxes.astype(int)
- boxes=boxes.astype(int)
- compare_o=np.asarray([boxes == boxes],dtype=int)
- compare=np.asarray([decoded_boxes[j] == boxes[j] for j in range(len(boxes))],dtype=int)
- print('Check 1: decoded pc level: ',pc_level)
- print('Check 2: decoded block level', departition_level)
- print('Check 3: decoded binstr ', binstr == decoded_binstr)
- print('Check 4: decoded boxes' ,np.count_nonzero(compare),compare.all())
- print('original boxes: ',np.count_nonzero(compare_o))
- # print('diff counts:',np.count_nonzero(compare_o)-np.count_nonzero(compare))
- # if np.count_nonzero(compare_o)-np.count_nonzero(compare) > 0:
- # for i in range(boxes.shape[0]):
- # compare=np.asarray([decoded_boxes[i] == boxes[i]],dtype=int)
- # print("block ",i,": ",compare.all())
- print('Decoding time: ', decoding_time)
-
- def decode_as_one(box, dec, voxelCNN, lock_inference):
- box_size = box.shape[1]
- #idx = int(np.log2(64 / box_size))
- '''idx = 0
- try:
- Model = voxelCNN[idx]
- except:
- print('Error! index of selecting model: ', idx, 'box shape: ', box.shape)
- '''
- lock_inference.acquire()
- probs = tl.ops.softmax(voxelCNN(box)[0, :, :, :, :], axis=-1)
- #probs = tf.nn.softmax(Model(box)[0, :, :, :, :], axis=-1)
- lock_inference.release()
- probs = probs[0:box_size, 0:box_size, 0:box_size, :]
- probs = np.asarray(probs, dtype='float32')
- start = time.time()
- decoded_box=np.zeros((box_size,box_size,box_size,1))
- end = time.time()
- for d in range(box_size):
- for h in range(box_size):
- for w in range(box_size):
- fre = [probs[d, h, w, 0], probs[d, h, w, 1], 0.]
- fre = np.asarray(fre)
- fre = (2 ** 10 * fre)
- fre = fre.astype(int)
- fre += 1
- freq = arithmetic_coding.NeuralFrequencyTable(fre)
- symbol = dec.read(freq)
- decoded_box[ d, h, w, 0] = symbol
- symbol = dec.read(freq)
- return dec, decoded_box
-
- def decompress_from_adaptive_freqs(boxes,heatmap, model_path64, output_path):
- no_box=len(boxes)
- bbox_max=boxes[0].shape[0] #64
- decoded_boxes = multiprocessing.Manager().list() #用数组的话只能展平为一维的,换算比较麻烦
- for i in range(no_box):
- decoded_boxes.append([])
- #decoded_boxes=np.zeros((no_box,bbox_max,bbox_max,bbox_max,1))
- lock = multiprocessing.Manager().Lock()
- lock_inference = multiprocessing.Manager().Lock()
- pool = multiprocessing.Pool(processes = MAX_PROCESSES) #维持执行的进程总数为processes,当一个进程执行完毕后会添加新的进程进去
-
- #for i in range(no_box):
- for i in range(no_box):##chang 1 to no_box if want to decode the all boxes
- print('Block ', i, '/', no_box, end='\r')
- box = []
- box.append(boxes[i, :, :, :, :])
- box = np.asarray(box)
- #print('number of non empty voxels: ', np.sum(box))
- curr_box_flag=heatmap[i][2]
- pool.apply_async(voxelCNN_decoding_thread, (i, box, output_path, model_path64, curr_box_flag, decoded_boxes, lock, lock_inference)) #非阻塞
-
- #print(' flag:', curr_box_flag)
- #dec,decoded_box,_=decoding_child_box_worker(box,voxelCNN,dec,curr_box_flag,idx)
- pool.close()
- pool.join() #调用join之前,先调用close函数,否则会出错。执行完close后不会有新的进程加入到pool,join函数等待所有子进程结束
- decoded_boxes = np.array(decoded_boxes)
- return decoded_boxes
-
- def voxelCNN_decoding_thread(block_idx, box, output_path, model_path64, curr_box_flag, decoded_boxes, lock, lock_inference):
- if BACKEND == 'tensorflow':
- import tensorflow as tf
- gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
- for gpu in gpus:
- tf.config.experimental.set_memory_growth(gpu, True) #推荐用这个,最多用到了8074M/15109M显存
- print(multiprocessing.current_process().pid, multiprocessing.current_process().name)
- outputfile = output_path + '.blocks' + str(block_idx) + '.bin'
- voxelCNN64 = VoxelCNN(depth=64, height=64, width=64, n_channel=1, output_channel=2, residual_blocks=2,n_filters=64)
- voxelCNN64.init_build(tl.layers.Input(shape=(1, 64, 64, 64, 1))) #这一步要有,走一遍前向推理,把前面没填的in_channels这些参数补上,再执行build
- tl.files.load_and_assign_npz(name=model_path64, network=voxelCNN64)
- voxelCNN64.set_eval()
- with open(outputfile, "rb") as inp:
- bitin = arithmetic_coding.BitInputStream(inp)
- dec = arithmetic_coding.ArithmeticDecoder(32, bitin)
- idx = 0
- dec,decoded_box,_=decoding_child_box_worker(box,voxelCNN64,dec,curr_box_flag,idx, lock_inference)
- lock.acquire()
- decoded_boxes[block_idx]=decoded_box
- lock.release()
-
- def decoding_child_box_worker(box, voxelCNN, dec, flag,idx, lock_inference):
- box_size = box.shape[1]
- decoded_box = np.zeros(( box_size, box_size, box_size, 1))
- if flag[idx]==2: #继续划分
- idx+=1
- child_bbox_max = int(box.shape[1] / 2)
- for d in range(2):
- for h in range(2):
- for w in range(2):
- child_box=box[:, d * child_bbox_max:(d + 1) * child_bbox_max,
- h * child_bbox_max:(h + 1) * child_bbox_max,
- w * child_bbox_max:(w + 1) * child_bbox_max, :]
- if(flag[idx]==0): #子box为空
- decoded_box[ d * child_bbox_max:(d + 1) * child_bbox_max,
- h * child_bbox_max:(h + 1) * child_bbox_max,
- w * child_bbox_max:(w + 1) * child_bbox_max, :] = 0
- idx+=1
- elif(flag[idx]==1):
- dec, decoded_box[d * child_bbox_max:(d + 1) * child_bbox_max,
- h * child_bbox_max:(h + 1) * child_bbox_max,
- w * child_bbox_max:(w + 1) * child_bbox_max, :]=decode_as_one(child_box,dec,voxelCNN, lock_inference)
- idx+=1
- elif(flag[idx]==2):
- dec, decoded_box[ d * child_bbox_max:(d + 1) * child_bbox_max,
- h * child_bbox_max:(h + 1) * child_bbox_max,
- w * child_bbox_max:(w + 1) * child_bbox_max, :],idx=decoding_child_box_worker(child_box,voxelCNN,dec,flag,idx, lock_inference)
- elif (flag[idx]==1): #box整体编解码
- idx+=1
- dec,decoded_box = decode_as_one(box, dec, voxelCNN, lock_inference)
- return dec,decoded_box,idx
-
-
- if __name__ == "__main__":
- parser = argparse.ArgumentParser(description='Encoding octree')
- parser.add_argument("-level", '--octreedepth', type=int,
- default=10,
- help='depth of input octree to pass for encoder')
- parser.add_argument("-depth", '--partitioningdepth', type=int,
- default=3,
- help='max depth to partition block')
- parser.add_argument("-ply", '--plypath', type=str, help='path to input ply file')
- parser.add_argument("-output", '--outputpath', type=str, help='path to output files')
- parser.add_argument("-model64", '--modelpath64', type=str, help='path to input model 64 .h5 file')
- parser.add_argument("-model32", '--modelpath32', type=str, help='path to input model 32 .h5 file')
- parser.add_argument("-model16", '--modelpath16', type=str, help='path to input model 16 .h5 file')
- parser.add_argument("-model8", '--modelpath8', type=str, help='path to input model 8 .h5 file')
- parser.add_argument("-signaling", '--signaling', type=str, help='special character for the output')
- args = parser.parse_args()
- voxelCNN_decoding([args.octreedepth, args.plypath, args.outputpath, args.modelpath64, args.modelpath32, args.modelpath16,
- args.modelpath8,args.partitioningdepth,args.signaling])
|