-
Notifications
You must be signed in to change notification settings - Fork 3
/
octree_np.py
95 lines (85 loc) · 3.58 KB
/
octree_np.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import os
import numpy as np
from tqdm import tqdm
########
## A simple Python implementation of octree coding
## Only point clouds with coordinates in the unit sphere (does not contain boundary values (i.e., 0 and 1)) are supported
########
def encode(pc, resolution, depth):
'''
input: 0 <= depth
'''
pc = getDecodeFromPc(pc, resolution, depth)
bits_ls = [[]]
for i in range(depth+1):
bits_ls.append([])
def include_points(startX, startY, startZ, endX, endY, endZ):
if np.any((startX <= pc[:, 0]) * (pc[:, 0] <= endX) * (startY <= pc[:, 1]) * (pc[:, 1] <= endY) * (startZ <= pc[:, 2]) * (pc[:, 2] <= endZ)):
return True
else:
return False
def octree(startX, startY, startZ, currdepth):
curr_cube_reso = resolution / (2 ** currdepth)
if include_points(startX, startY, startZ, startX+curr_cube_reso, startY+curr_cube_reso, startZ+curr_cube_reso):
bits_ls[currdepth].append(1)
if currdepth == depth:
return
next_cube_reso = curr_cube_reso / 2
octree(startX, startY, startZ, currdepth+1)
octree(startX, startY, startZ+next_cube_reso, currdepth+1)
octree(startX, startY+next_cube_reso, startZ, currdepth+1)
octree(startX, startY+next_cube_reso, startZ+next_cube_reso, currdepth+1)
octree(startX+next_cube_reso, startY, startZ, currdepth+1)
octree(startX+next_cube_reso, startY, startZ+next_cube_reso, currdepth+1)
octree(startX+next_cube_reso, startY+next_cube_reso, startZ, currdepth+1)
octree(startX+next_cube_reso, startY+next_cube_reso, startZ+next_cube_reso, currdepth+1)
else:
bits_ls[currdepth].append(0)
octree(0,0,0,0)
bits = [i for ls in bits_ls for i in ls]
del bits[0]
bits = np.array(bits)
return bits
def decode(bits, resolution):
# build the bits_ls
bits_ls = [[1]]
bits = bits.tolist()
n = 8
while True:
bits_group = bits[:n]
del bits[:n]
bits_ls.append(bits_group)
n = sum(bits_group) * 8
if len(bits) == 0:
break
depth = len(bits_ls) - 1
pc = []
#print(depth)
def dec(startX, startY, startZ, currdepth):
curr_cube_reso = resolution / (2 ** currdepth)
b = bits_ls[currdepth].pop(0)
if b == 1:
#print(b)
if currdepth == depth:
pc.append([startX+curr_cube_reso/2, startY+curr_cube_reso/2, startZ+curr_cube_reso/2])
return
next_cube_reso = curr_cube_reso / 2
dec(startX, startY, startZ, currdepth+1)
dec(startX, startY, startZ+next_cube_reso, currdepth+1)
dec(startX, startY+next_cube_reso, startZ, currdepth+1)
dec(startX, startY+next_cube_reso, startZ+next_cube_reso, currdepth+1)
dec(startX+next_cube_reso, startY, startZ, currdepth+1)
dec(startX+next_cube_reso, startY, startZ+next_cube_reso, currdepth+1)
dec(startX+next_cube_reso, startY+next_cube_reso, startZ, currdepth+1)
dec(startX+next_cube_reso, startY+next_cube_reso, startZ+next_cube_reso, currdepth+1)
else:
return
dec(0,0,0,0)
pc = np.array(pc)
return pc
def getDecodeFromPc(pc, resolution, depth):
cube_reso = resolution / (2 ** depth)
pc_octree = (pc // cube_reso * cube_reso) + (cube_reso / 2)
pc_octree = np.unique(pc_octree, axis=0)
#print(pc_octree)
return pc_octree