def test_bsp(self, config):
if self.checkpoint_manager.latest_checkpoint:
self.checkpoint.restore(self.checkpoint_manager.latest_checkpoint)
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
return
w2 = self.bsp_network.generator.convex_layer_weights.numpy()
dima = self.test_size
dim = self.real_size
multiplier = int(dim/dima)
multiplier2 = multiplier*multiplier
for t in range(config.start, min(len(self.data_pixels),config.end)):
model_float = np.ones([self.real_size,self.real_size,self.real_size,self.c_dim],np.float32)
batch_view = self.data_pixels[t:t+1,self.test_idx].astype(np.float32)/255.0
_, out_m, _,_ = self.bsp_network(batch_view, None, None, None, is_training=False)
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
minib = i*multiplier2+j*multiplier+k
point_coord = self.coords[minib:minib+1]
_,_, model_out, _ = self.bsp_network(None, None, out_m, point_coord, is_training=False)
model_float[self.aux_x+i,self.aux_y+j,self.aux_z+k,:] = np.reshape(model_out, [self.test_size,self.test_size,self.test_size,self.c_dim])
out_m = out_m.numpy()
bsp_convex_list = []
model_float = model_float<0.01
model_float_sum = np.sum(model_float,axis=3)
for i in range(self.c_dim):
slice_i = model_float[:,:,:,i]
if np.max(slice_i)>0: #if one voxel is inside a convex
if np.min(model_float_sum-slice_i*2)>=0: #if this convex is redundant, i.e. the convex is inside the shape
model_float_sum = model_float_sum-slice_i
else:
box = []
for j in range(self.p_dim):
if w2[j,i]>0.01:
a = -out_m[0,0,j]
b = -out_m[0,1,j]
c = -out_m[0,2,j]
d = -out_m[0,3,j]
box.append([a,b,c,d])
if len(box)>0:
bsp_convex_list.append(np.array(box,np.float32))
#print(bsp_convex_list)
print(len(bsp_convex_list))
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(config.sample_dir+"/"+str(t)+"_bsp.ply", vertices, polygons)
def reconstruct_object(image_path, checkpoint_path, output_path):
device = torch.device('cpu:0')
real_size = 64 # Output point-value voxel grid size in testing.
test_size = 32 # Related to testing batch_size.
test_point_batch_size = test_size * test_size * test_size
# Get coords.
dima = test_size
dim = real_size
aux_x = np.zeros([dima, dima, dima], np.uint8)
aux_y = np.zeros([dima, dima, dima], np.uint8)
aux_z = np.zeros([dima, dima, dima], np.uint8)
multiplier = int(dim / dima)
multiplier2 = multiplier * multiplier
multiplier3 = multiplier * multiplier * multiplier
for i in range(dima):
for j in range(dima):
for k in range(dima):
aux_x[i, j, k] = i * multiplier
aux_y[i, j, k] = j * multiplier
aux_z[i, j, k] = k * multiplier
coords = np.zeros([multiplier3, dima, dima, dima, 3], np.float32)
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
coords[i * multiplier2 + j * multiplier + k, :, :, :,
0] = aux_x + i
coords[i * multiplier2 + j * multiplier + k, :, :, :,
1] = aux_y + j
coords[i * multiplier2 + j * multiplier + k, :, :, :,
2] = aux_z + k
coords = (coords + 0.5) / dim - 0.5
coords = np.reshape(coords, [multiplier3, test_point_batch_size, 3])
coords = np.concatenate(
[coords,
np.ones([multiplier3, test_point_batch_size, 1], np.float32)],
axis=2)
coords = torch.from_numpy(coords)
coords = coords.to(device)
# Instantiate model.
ef_dim = 32
p_dim = 4096 # Number of planes.
c_dim = 256 # Number of convexes.
print("Instantiating model...")
bsp_net = model.BSPNet(ef_dim,
p_dim,
c_dim,
img_ef_dim=64,
z_dim=ef_dim * 8)
bsp_net.to(device)
bsp_net.eval() # The network must be set to evaluation mode.
# Load pre-trained model.
print("Loading model weights loaded from checkpoint...")
bsp_net.load_state_dict(torch.load(checkpoint_path))
# Load the (input) image and convert it to tensor.
print("Loading image...")
image = PIL.Image.open(image_path)
image = IMAGE_TRANSFORM(image)
# BSP-Net will determine:
# 1. which planes to use
# 2. how they are used to form convexes
# 3. how convexes form the object
w2 = bsp_net.generator.convex_layer_weights.detach().cpu().numpy()
multiplier = int(real_size / test_size)
multiplier2 = multiplier * multiplier
model_float = np.ones([real_size, real_size, real_size, c_dim], np.float32)
model_float_combined = np.ones([real_size, real_size, real_size],
np.float32)
batch_view = torch.reshape(image, shape=(1, 1, 128, 128))
batch_view = batch_view.to(device)
print("Predicting planes from image...")
_, plane_m, _, _ = bsp_net(batch_view, None, None, None, is_training=False)
print("Predicting partitions...")
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
minib = i * multiplier2 + j * multiplier + k
point_coord = coords[minib:minib + 1]
_, _, model_out, model_out_combined = bsp_net(
None, None, plane_m, point_coord, is_training=False)
model_float[aux_x + i, aux_y + j, aux_z + k, :] = np.reshape(
model_out.detach().cpu().numpy(),
[test_size, test_size, test_size, c_dim])
model_float_combined[
aux_x + i, aux_y + j, aux_z + k] = np.reshape(
model_out_combined.detach().cpu().numpy(),
[test_size, test_size, test_size])
# Gather convexes from planes + partitions.
bsp_convex_list = []
plane_m = plane_m.detach().cpu().numpy()
model_float = model_float < 0.01
model_float_sum = np.sum(model_float, axis=3)
for i in range(c_dim):
slice_i = model_float[:, :, :, i]
if np.max(slice_i) > 0:
box = []
for j in range(p_dim):
if w2[j, i] > 0.01:
a = -plane_m[0, 0, j]
b = -plane_m[0, 1, j]
c = -plane_m[0, 2, j]
d = -plane_m[0, 3, j]
box.append([a, b, c, d])
if len(box) > 0:
bsp_convex_list.append(np.array(box, np.float32))
# Convert bspt to mesh.
print("Converting list of convexes to watertight mesh...")
vertices, polygons = bspt.get_mesh_watertight(bsp_convex_list)
# Convert mesh .ply to .obj
# Then save mesh on disk.
print("Saving mesh on disk...")
outdir_path = os.path.dirname(output_path)
os.makedirs(outdir_path, exist_ok=True)
utils.write_ply_polygon(output_path, vertices, polygons)
def test_mesh_point(self, config):
#load previous checkpoint
if not self.load(): exit(-1)
w2 = self.bsp_network.generator.convex_layer_weights.detach().cpu(
).numpy()
dima = self.test_size
dim = self.real_size
multiplier = int(dim / dima)
multiplier2 = multiplier * multiplier
self.bsp_network.eval()
for t in range(config.start, min(len(self.data_pixels), config.end)):
print(t)
model_float = np.ones(
[self.real_size, self.real_size, self.real_size, self.c_dim],
np.float32)
model_float_combined = np.ones(
[self.real_size, self.real_size, self.real_size], np.float32)
batch_view = self.data_pixels[t:t + 1, self.test_idx].astype(
np.float32) / 255.0
batch_view = torch.from_numpy(batch_view)
batch_view = batch_view.to(self.device)
_, out_m, _, _ = self.bsp_network(batch_view,
None,
None,
None,
is_training=False)
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
minib = i * multiplier2 + j * multiplier + k
point_coord = self.coords[minib:minib + 1]
_, _, model_out, model_out_combined = self.bsp_network(
None, None, out_m, point_coord, is_training=False)
model_float[self.aux_x + i, self.aux_y + j,
self.aux_z + k, :] = np.reshape(
model_out.detach().cpu().numpy(), [
self.test_size, self.test_size,
self.test_size, self.c_dim
])
model_float_combined[
self.aux_x + i, self.aux_y + j,
self.aux_z + k] = np.reshape(
model_out_combined.detach().cpu().numpy(), [
self.test_size, self.test_size,
self.test_size
])
out_m_ = out_m.detach().cpu().numpy()
# whether to use post processing to remove convexes that are inside the shape
post_processing_flag = False
if post_processing_flag:
bsp_convex_list = []
model_float = model_float < 0.01
model_float_sum = np.sum(model_float, axis=3)
unused_convex = np.ones([self.c_dim], np.float32)
for i in range(self.c_dim):
slice_i = model_float[:, :, :, i]
if np.max(slice_i) > 0: #if one voxel is inside a convex
if np.min(
model_float_sum - slice_i * 2
) >= 0: #if this convex is redundant, i.e. the convex is inside the shape
model_float_sum = model_float_sum - slice_i
else:
box = []
for j in range(self.p_dim):
if w2[j, i] > 0.01:
a = -out_m_[0, 0, j]
b = -out_m_[0, 1, j]
c = -out_m_[0, 2, j]
d = -out_m_[0, 3, j]
box.append([a, b, c, d])
if len(box) > 0:
bsp_convex_list.append(
np.array(box, np.float32))
unused_convex[i] = 0
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(
config.sample_dir + "/" + str(t) + "_bsp.ply", vertices,
polygons)
#output obj
#write_obj_polygon(config.sample_dir+"/"+str(t)+"_bsp.obj", vertices, polygons)
#sample surface points
sampled_points_normals = sample_points_polygon(
vertices, polygons, 16384)
#check point inside shape or not
point_coord = np.reshape(
sampled_points_normals[:, :3] +
sampled_points_normals[:, 3:] * 1e-4, [1, -1, 3])
point_coord = np.concatenate([
point_coord,
np.ones([1, point_coord.shape[1], 1], np.float32)
],
axis=2)
_, _, _, sample_points_value = self.bsp_network(
None,
None,
out_m,
torch.from_numpy(point_coord).to(self.device),
convex_mask=torch.from_numpy(
np.reshape(unused_convex, [1, 1, -1])).to(self.device),
is_training=False)
sample_points_value = sample_points_value.detach().cpu().numpy(
)
sampled_points_normals = sampled_points_normals[
sample_points_value[0, :, 0] > 1e-4]
print(len(bsp_convex_list), len(sampled_points_normals))
np.random.shuffle(sampled_points_normals)
write_ply_point_normal(
config.sample_dir + "/" + str(t) + "_pc.ply",
sampled_points_normals[:4096])
else:
bsp_convex_list = []
model_float = model_float < 0.01
model_float_sum = np.sum(model_float, axis=3)
for i in range(self.c_dim):
slice_i = model_float[:, :, :, i]
if np.max(slice_i) > 0: #if one voxel is inside a convex
#if np.min(model_float_sum-slice_i*2)>=0: #if this convex is redundant, i.e. the convex is inside the shape
# model_float_sum = model_float_sum-slice_i
#else:
box = []
for j in range(self.p_dim):
if w2[j, i] > 0.01:
a = -out_m_[0, 0, j]
b = -out_m_[0, 1, j]
c = -out_m_[0, 2, j]
d = -out_m_[0, 3, j]
box.append([a, b, c, d])
if len(box) > 0:
bsp_convex_list.append(np.array(box, np.float32))
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(
config.sample_dir + "/" + str(t) + "_bsp.ply", vertices,
polygons)
#output obj
#write_obj_polygon(config.sample_dir+"/"+str(t)+"_bsp.obj", vertices, polygons)
#sample surface points
sampled_points_normals = sample_points_polygon_vox64(
vertices, polygons, model_float_combined, 16384)
#check point inside shape or not
point_coord = np.reshape(
sampled_points_normals[:, :3] +
sampled_points_normals[:, 3:] * 1e-4, [1, -1, 3])
point_coord = np.concatenate([
point_coord,
np.ones([1, point_coord.shape[1], 1], np.float32)
],
axis=2)
_, _, _, sample_points_value = self.bsp_network(
None,
None,
out_m,
torch.from_numpy(point_coord).to(self.device),
is_training=False)
sample_points_value = sample_points_value.detach().cpu().numpy(
)
sampled_points_normals = sampled_points_normals[
sample_points_value[0, :, 0] > 1e-4]
print(len(bsp_convex_list), len(sampled_points_normals))
np.random.shuffle(sampled_points_normals)
write_ply_point_normal(
config.sample_dir + "/" + str(t) + "_pc.ply",
sampled_points_normals[:4096])
def test_mesh_point(self, config):
could_load, checkpoint_counter = self.load(self.checkpoint_dir)
if could_load:
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
return
w2 = self.sess.run(self.cw2, feed_dict={})
dima = self.test_size
dim = self.real_size
multiplier = int(dim/dima)
multiplier2 = multiplier*multiplier
for t in range(config.start, min(len(self.data_voxels),config.end)):
print(t)
model_float = np.ones([self.real_size,self.real_size,self.real_size,self.c_dim],np.float32)
model_float_combined = np.ones([self.real_size,self.real_size,self.real_size],np.float32)
batch_voxels = self.data_voxels[t:t+1]
out_m, out_b = self.sess.run([self.sE_m, self.sE_b],
feed_dict={
self.vox3d: batch_voxels,
})
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
minib = i*multiplier2+j*multiplier+k
model_out, model_out_combined = self.sess.run([self.zG2, self.zG],
feed_dict={
self.plane_m: out_m,
self.plane_b: out_b,
self.point_coord: self.coords[minib:minib+1],
})
model_float[self.aux_x+i,self.aux_y+j,self.aux_z+k,:] = np.reshape(model_out, [self.test_size,self.test_size,self.test_size,self.c_dim])
model_float_combined[self.aux_x+i,self.aux_y+j,self.aux_z+k] = np.reshape(model_out_combined, [self.test_size,self.test_size,self.test_size])
# whether to use post processing to remove convexes that are inside the shape
post_processing_flag = False
if post_processing_flag:
bsp_convex_list = []
model_float = model_float<0.01
model_float_sum = np.sum(model_float,axis=3)
unused_convex = np.ones([self.c_dim], np.float32)
for i in range(self.c_dim):
slice_i = model_float[:,:,:,i]
if np.max(slice_i)>0: #if one voxel is inside a convex
if np.min(model_float_sum-slice_i*2)>=0: #if this convex is redundant, i.e. the convex is inside the shape
model_float_sum = model_float_sum-slice_i
else:
box = []
for j in range(self.p_dim):
if w2[j,i]>0.01:
a = -out_m[0,0,j]
b = -out_m[0,1,j]
c = -out_m[0,2,j]
d = -out_b[0,0,j]
box.append([a,b,c,d])
if len(box)>0:
bsp_convex_list.append(np.array(box,np.float32))
unused_convex[i] = 0
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(config.sample_dir+"/"+str(t)+"_bsp.ply", vertices, polygons)
#output obj
#write_obj_polygon(config.sample_dir+"/"+str(t)+"_bsp.obj", vertices, polygons)
#sample surface points
sampled_points_normals = sample_points_polygon(vertices, polygons, 16384)
#check point inside shape or not
sample_points_value = self.sess.run(self.zmG,
feed_dict={
self.plane_m: out_m,
self.plane_b: out_b,
self.convex_mask: np.reshape(unused_convex, [1,1,-1]),
self.point_coord: np.reshape(sampled_points_normals[:,:3]+sampled_points_normals[:,3:]*1e-4, [1,-1,3]),
})
sampled_points_normals = sampled_points_normals[sample_points_value[0,:,0]>1e-4]
print(len(bsp_convex_list), len(sampled_points_normals))
np.random.shuffle(sampled_points_normals)
write_ply_point_normal(config.sample_dir+"/"+str(t)+"_pc.ply", sampled_points_normals[:4096])
else:
bsp_convex_list = []
model_float = model_float<0.01
model_float_sum = np.sum(model_float,axis=3)
for i in range(self.c_dim):
slice_i = model_float[:,:,:,i]
if np.max(slice_i)>0: #if one voxel is inside a convex
#if np.min(model_float_sum-slice_i*2)>=0: #if this convex is redundant, i.e. the convex is inside the shape
# model_float_sum = model_float_sum-slice_i
#else:
box = []
for j in range(self.p_dim):
if w2[j,i]>0.01:
a = -out_m[0,0,j]
b = -out_m[0,1,j]
c = -out_m[0,2,j]
d = -out_b[0,0,j]
box.append([a,b,c,d])
if len(box)>0:
bsp_convex_list.append(np.array(box,np.float32))
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(config.sample_dir+"/"+str(t)+"_bsp.ply", vertices, polygons)
#output obj
#write_obj_polygon(config.sample_dir+"/"+str(t)+"_bsp.obj", vertices, polygons)
#sample surface points
sampled_points_normals = sample_points_polygon_vox64(vertices, polygons, model_float_combined, 16384)
#check point inside shape or not
sample_points_value = self.sess.run(self.zG,
feed_dict={
self.plane_m: out_m,
self.plane_b: out_b,
self.point_coord: np.reshape(sampled_points_normals[:,:3]+sampled_points_normals[:,3:]*1e-4, [1,-1,3]),
})
sampled_points_normals = sampled_points_normals[sample_points_value[0,:,0]>1e-4]
print(len(bsp_convex_list), len(sampled_points_normals))
np.random.shuffle(sampled_points_normals)
write_ply_point_normal(config.sample_dir+"/"+str(t)+"_pc.ply", sampled_points_normals[:4096])
def test_bsp(self, config):
could_load, checkpoint_counter = self.load(self.checkpoint_dir)
if could_load:
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
return
w2 = self.sess.run(self.cw2, feed_dict={})
dima = self.test_size
dim = self.real_size
multiplier = int(dim/dima)
multiplier2 = multiplier*multiplier
for t in range(config.start, min(len(self.data_voxels),config.end)):
model_float = np.ones([self.real_size,self.real_size,self.real_size,self.c_dim],np.float32)
batch_voxels = self.data_voxels[t:t+1]
out_m, out_b = self.sess.run([self.sE_m, self.sE_b],
feed_dict={
self.vox3d: batch_voxels,
})
for i in range(multiplier):
for j in range(multiplier):
for k in range(multiplier):
minib = i*multiplier2+j*multiplier+k
model_out = self.sess.run(self.zG2,
feed_dict={
self.plane_m: out_m,
self.plane_b: out_b,
self.point_coord: self.coords[minib:minib+1],
})
model_float[self.aux_x+i,self.aux_y+j,self.aux_z+k,:] = np.reshape(model_out, [self.test_size,self.test_size,self.test_size,self.c_dim])
bsp_convex_list = []
model_float = model_float<0.01
model_float_sum = np.sum(model_float,axis=3)
for i in range(self.c_dim):
slice_i = model_float[:,:,:,i]
if np.max(slice_i)>0: #if one voxel is inside a convex
if np.min(model_float_sum-slice_i*2)>=0: #if this convex is redundant, i.e. the convex is inside the shape
model_float_sum = model_float_sum-slice_i
else:
box = []
for j in range(self.p_dim):
if w2[j,i]>0.01:
a = -out_m[0,0,j]
b = -out_m[0,1,j]
c = -out_m[0,2,j]
d = -out_b[0,0,j]
box.append([a,b,c,d])
if len(box)>0:
bsp_convex_list.append(np.array(box,np.float32))
#print(bsp_convex_list)
print(len(bsp_convex_list))
#convert bspt to mesh
#vertices, polygons = get_mesh(bsp_convex_list)
#use the following alternative to merge nearby vertices to get watertight meshes
vertices, polygons = get_mesh_watertight(bsp_convex_list)
#output ply
write_ply_polygon(config.sample_dir+"/"+str(t)+"_bsp.ply", vertices, polygons)