def depth2sphere_tool(depth_path, mask_path):
proj_depth = Camera_back_projection_layer().cuda()
render_spherical = spherical_proj.render_spherical().cuda()
depth = skimage.io.imread(depth_path, as_gray=True)
mask = skimage.io.imread(mask_path, as_gray=True)
mine = False
if mine:
with open("data/0006.json", "r") as j_file: # DRILL
data = json.load(j_file)[996]
# with open("data/0000.json", "r") as j_file: # Box
# data = json.load(j_file)[0]
#
depth_minmax = data["depth_minmax"]
else:
depth_minmax = np.load(
'data/02691156_fff513f407e00e85a9ced22d91ad7027_view019.npy')
print(depth_minmax)
print(np.amin(depth))
print(np.amax(depth))
pred_abs_depth = get_abs_depth(depth, mask, depth_minmax)
print(type(pred_abs_depth))
pred_abs_depth = pred_abs_depth.type(torch.Tensor).cuda()
#pred_abs_depth = pred_abs_depth.squeeze().cpu().numpy()
print(depth_minmax)
# print(np.amin(pred_abs_depth))
# print(np.amax(pred_abs_depth))
# plt.imshow(pred_abs_depth)
# plt.show()
proj = proj_depth(pred_abs_depth)
proj_df = proj.squeeze()
voxel = proj_df.cpu().numpy()
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.set_title('partial_voxel')
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_aspect('equal')
ax.voxels(voxel[::1, ::1, ::1], edgecolor="k")
ax.view_init(0, 180)
plt.draw()
plt.savefig('vox_partial_sampled.png')
plt.show()
plt.clf()
print(proj.shape)
sph = render_spherical(torch.clamp(proj * 50, 1e-5, 1 - 1e-5))
sph = sph.squeeze().cpu().numpy()
# sph = np.flipud(sph)
print(sph.shape)
# plt.ion()
plt.imshow(sph)
plt.savefig('depth2sph.png')
def __init__(self, opt, base_class=Model):
super().__init__()
self.net1 = Net1([3, 1, 1], ['normal', 'depth', 'silhou'],
pred_depth_minmax=True)
self.net2 = Uresnet([1], ['spherical'], input_planes=1)
self.base_class = base_class
self.proj_depth = Camera_back_projection_layer()
self.render_spherical = render_spherical()
self.joint_train = opt.joint_train
self.load_offline = opt.load_offline
self.padding_margin = opt.padding_margin
if opt.net1_path:
state_dicts = torch.load(opt.net1_path)
self.net1.load_state_dict(state_dicts['nets'][0])
def vox2sphere(invoxel=None):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if invoxel is None:
mine = True
if mine:
# with open('data/rotated_mesh.binvox', 'rb') as f:
# with open('data/incomplete_mesh.binvox', 'rb') as f:
# m1 = binvox_rw.read_as_3d_array(f)
#
# voxel = m1.data.astype(np.float32)
box = True
if box:
data = np.load(
"/mnt/Extra/dsl/GenRe-ShapeHD/gen_sphr/data/rotated_mesh_voxelized.npz"
)
else:
data = np.load(
"/mnt/Extra/dsl/GenRe-ShapeHD/gen_sphr/data/rotated_mesh_voxelized_drill.npz"
)
voxel = data['voxel'].astype(np.float64)
else:
data = np.load(
'data/02691156_fff513f407e00e85a9ced22d91ad7027_view019_gt_rotvox_samescale_128.npz'
)
voxel = data['voxel']
# plot_voxel(voxel)
else:
voxel = invoxel
projector = spherical_proj.render_spherical().to(device)
voxeltensor = torch.from_numpy(voxel).float().to(device)
voxeltensor = voxeltensor.unsqueeze(0)
voxeltensor = voxeltensor.unsqueeze(0)
spherical = projector.forward(voxeltensor)
spherical = spherical.squeeze().to("cpu").numpy()
spherical = np.flipud(spherical)
# spherical = np.fliplr(spherical)
plt.imshow(spherical)
plt.savefig('vox2sphere.png'
) # if invoxel is None else plt.savefig('depth2sphere.png')
# plt.show()
sphere2vox(spherical)
def vox2sphere(data):
# vox_path = data['mask'].replace('mask', 'voxel').replace('png', 'npz').replace('media', 'mnt/hammin')
# if not os.path.exists(vox_path): return None
# voxel = np.load(vox_path)['voxel'].astype(np.float32)
voxel = data["voxel"].astype(np.float32)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# if invoxel is None:
# mine = True
# if mine:
# # with open('data/rotated_mesh.binvox', 'rb') as f:
# # with open('data/incomplete_mesh.binvox', 'rb') as f:
# # m1 = binvox_rw.read_as_3d_array(f)
# #
# # voxel = m1.data.astype(np.float32)
# box = True
# if box:
# data = np.load("/mnt/Extra/dsl/GenRe-ShapeHD/gen_sphr/data/rotated_mesh_voxelized.npz")
# else:
# data = np.load("/mnt/Extra/dsl/GenRe-ShapeHD/gen_sphr/data/rotated_mesh_voxelized_drill.npz")
# voxel = data['voxel'].astype(np.float64)
#
# else:
# data = np.load(
# 'data/02691156_fff513f407e00e85a9ced22d91ad7027_view019_gt_rotvox_samescale_128.npz')
# voxel = data['voxel']
# # plot_voxel(voxel)
#
# else:
# voxel = invoxel
projector = spherical_proj.render_spherical().to(device)
voxeltensor = torch.from_numpy(voxel).float().to(device)
voxeltensor = voxeltensor.unsqueeze(0)
voxeltensor = voxeltensor.unsqueeze(0)
spherical = projector.forward(voxeltensor)
spherical = spherical.squeeze().to("cpu").numpy()
spherical = np.flipud(spherical)
return spherical
# The videos to process
folder_names = os.listdir(input_dir)
videos = []
for folder in folder_names:
if os.path.isdir(input_dir + '/' + folder):
videos.append(folder)
videos.sort()
# set the visualizer
my_visualizer = Visualizer(n_workers=getattr(opt, 'vis_workers', 4),
param_f=getattr(opt, 'vis_param_f', None))
# set the models to use
# first, set the Spherical map rendering
render_spherical_map = render_spherical().to(device)
print('>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
print('The model for rendering spherical maps has been set successfully')
# second, set the genre full model. In this program, use the second model which
# is spherical map inpainting net, to inpaint the spherical map, and use the refine
# model which refines the voxels to get the final shape
Full_model = models.get_model(opt.net, test=True)
full_model = Full_model(opt, logger)
full_model.to(device)
full_model.eval()
print('>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
print('The model for inpainting sphercial maps has been set successfully')
print('>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
print('The model for refining voxels has been set successfully')
# Process for each video