126 lines
5.7 KiB
Python
126 lines
5.7 KiB
Python
|
|
import os
|
|
import json
|
|
import time
|
|
import subprocess
|
|
import tempfile
|
|
import shutil
|
|
import numpy as np
|
|
from utils.data_load import DataLoadUtil
|
|
from utils.reconstruction import ReconstructionUtil
|
|
from utils.pts import PtsUtil
|
|
class RenderUtil:
|
|
target_mask_label = (0, 255, 0)
|
|
display_table_mask_label = (0, 0, 255)
|
|
random_downsample_N = 32768
|
|
min_z = 0.2
|
|
max_z = 0.5
|
|
|
|
@staticmethod
|
|
def get_world_points_and_normal(depth, mask, normal, cam_intrinsic, cam_extrinsic, random_downsample_N):
|
|
z = depth[mask]
|
|
i, j = np.nonzero(mask)
|
|
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
|
|
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
|
|
|
|
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
|
|
normal_camera = normal[mask].reshape(-1, 3)
|
|
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
|
|
points_camera, random_downsample_N, require_idx=True
|
|
)
|
|
if len(sampled_target_points) == 0:
|
|
return np.zeros((0, 3)), np.zeros((0, 3))
|
|
sampled_normal_camera = normal_camera[idx]
|
|
|
|
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
|
|
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
|
|
|
|
return points_camera_world, sampled_normal_camera
|
|
|
|
@staticmethod
|
|
def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic, random_downsample_N):
|
|
z = depth[mask]
|
|
i, j = np.nonzero(mask)
|
|
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
|
|
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
|
|
|
|
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
|
|
sampled_target_points = PtsUtil.random_downsample_point_cloud(
|
|
points_camera, random_downsample_N
|
|
)
|
|
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
|
|
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
|
|
|
|
return points_camera_world
|
|
|
|
@staticmethod
|
|
def render_pts(cam_pose, scene_path, script_path, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False):
|
|
|
|
nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path)
|
|
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
params = {
|
|
"cam_pose": nO_to_world_pose.tolist(),
|
|
"scene_path": scene_path
|
|
}
|
|
scene_info_path = os.path.join(scene_path, "scene_info.json")
|
|
shutil.copy(scene_info_path, os.path.join(temp_dir, "scene_info.json"))
|
|
params_data_path = os.path.join(temp_dir, "params.json")
|
|
with open(params_data_path, 'w') as f:
|
|
json.dump(params, f)
|
|
start_time = time.time()
|
|
result = subprocess.run([
|
|
'blender', '-b', '-P', script_path, '--', temp_dir
|
|
], capture_output=True, text=True)
|
|
end_time = time.time()
|
|
print(result)
|
|
print(f"-- Time taken for blender: {end_time - start_time} seconds")
|
|
if result.returncode != 0:
|
|
print("Blender script failed:")
|
|
print(result.stderr)
|
|
return None
|
|
path = os.path.join(temp_dir, "tmp")
|
|
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
|
|
depth_L, depth_R = DataLoadUtil.load_depth(
|
|
path, cam_info["near_plane"],
|
|
cam_info["far_plane"],
|
|
binocular=True
|
|
)
|
|
start_time = time.time()
|
|
mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
|
|
normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
|
|
''' target points '''
|
|
mask_img_L = mask_L
|
|
mask_img_R = mask_R
|
|
|
|
target_mask_img_L = (mask_L == RenderUtil.target_mask_label).all(axis=-1)
|
|
target_mask_img_R = (mask_R == RenderUtil.target_mask_label).all(axis=-1)
|
|
|
|
|
|
sampled_target_points_L, sampled_target_normal_L = RenderUtil.get_world_points_and_normal(depth_L,target_mask_img_L,normal_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], RenderUtil.random_downsample_N)
|
|
sampled_target_points_R = RenderUtil.get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], RenderUtil.random_downsample_N )
|
|
|
|
|
|
has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
|
|
if has_points:
|
|
target_points, overlap_idx = PtsUtil.get_overlapping_points(
|
|
sampled_target_points_L, sampled_target_points_R, voxel_threshold, require_idx=True
|
|
)
|
|
sampled_target_normal_L = sampled_target_normal_L[overlap_idx]
|
|
|
|
if has_points:
|
|
has_points = target_points.shape[0] > 0
|
|
|
|
if has_points:
|
|
target_points, target_normals = PtsUtil.filter_points(
|
|
target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(RenderUtil.min_z, RenderUtil.max_z)
|
|
)
|
|
|
|
if not has_points:
|
|
target_points = np.zeros((0, 3))
|
|
target_normals = np.zeros((0, 3))
|
|
end_time = time.time()
|
|
print(f"-- Time taken for processing: {end_time - start_time} seconds")
|
|
#import ipdb; ipdb.set_trace()
|
|
return target_points, target_normals |