add preprocess
This commit is contained in:
hofee
@@ -157,8 +157,8 @@ class DataLoadUtil:
|
||||
return depth_meters
|
||||
|
||||
@staticmethod
|
||||
def load_seg(path, binocular=False):
|
||||
if binocular:
|
||||
def load_seg(path, binocular=False, left_only=False):
|
||||
if binocular and not left_only:
|
||||
|
||||
def clean_mask(mask_image):
|
||||
green = [0, 255, 0, 255]
|
||||
@@ -182,11 +182,41 @@ class DataLoadUtil:
|
||||
mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED))
|
||||
return mask_image_L, mask_image_R
|
||||
else:
|
||||
mask_path = os.path.join(
|
||||
os.path.dirname(path), "mask", os.path.basename(path) + ".png"
|
||||
)
|
||||
if binocular and left_only:
|
||||
mask_path = os.path.join(
|
||||
os.path.dirname(path), "mask", os.path.basename(path) + "_L.png"
|
||||
)
|
||||
else:
|
||||
mask_path = os.path.join(
|
||||
os.path.dirname(path), "mask", os.path.basename(path) + ".png"
|
||||
)
|
||||
mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
|
||||
return mask_image
|
||||
|
||||
@staticmethod
|
||||
def load_normal(path, binocular=False, left_only=False):
|
||||
if binocular and not left_only:
|
||||
normal_path_L = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
|
||||
)
|
||||
normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_UNCHANGED)
|
||||
normal_path_R = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + "_R.png"
|
||||
)
|
||||
normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_UNCHANGED)
|
||||
return normal_image_L, normal_image_R
|
||||
else:
|
||||
if binocular and left_only:
|
||||
normal_path = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
|
||||
)
|
||||
|
||||
else:
|
||||
normal_path = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + ".png"
|
||||
)
|
||||
normal_image = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
|
||||
return normal_image
|
||||
|
||||
@staticmethod
|
||||
def load_label(path):
|
||||
@@ -273,7 +303,7 @@ class DataLoadUtil:
|
||||
|
||||
@staticmethod
|
||||
def get_target_point_cloud(
|
||||
depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0, 255, 0, 255)
|
||||
depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0, 255, 0, 255), require_full_points=False
|
||||
):
|
||||
h, w = depth.shape
|
||||
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
|
||||
@@ -293,10 +323,11 @@ class DataLoadUtil:
|
||||
)
|
||||
|
||||
target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
|
||||
return {
|
||||
data = {
|
||||
"points_world": target_points_world,
|
||||
"points_camera": target_points_camera,
|
||||
}
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def get_point_cloud(depth, cam_intrinsic, cam_extrinsic):
|
||||
@@ -323,7 +354,8 @@ class DataLoadUtil:
|
||||
voxel_size=0.005,
|
||||
target_mask_label=(0, 255, 0, 255),
|
||||
display_table_mask_label=(0, 0, 255, 255),
|
||||
get_display_table_pts=False
|
||||
get_display_table_pts=False,
|
||||
require_normal=False,
|
||||
):
|
||||
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
|
||||
if binocular:
|
||||
@@ -351,34 +383,9 @@ class DataLoadUtil:
|
||||
point_cloud_R = PtsUtil.random_downsample_point_cloud(
|
||||
point_cloud_R, random_downsample_N
|
||||
)
|
||||
overlap_points = DataLoadUtil.get_overlapping_points(
|
||||
overlap_points = PtsUtil.get_overlapping_points(
|
||||
point_cloud_L, point_cloud_R, voxel_size
|
||||
)
|
||||
if get_display_table_pts:
|
||||
display_pts_L = DataLoadUtil.get_target_point_cloud(
|
||||
depth_L,
|
||||
cam_info["cam_intrinsic"],
|
||||
cam_info["cam_to_world"],
|
||||
mask_L,
|
||||
display_table_mask_label,
|
||||
)["points_world"]
|
||||
display_pts_R = DataLoadUtil.get_target_point_cloud(
|
||||
depth_R,
|
||||
cam_info["cam_intrinsic"],
|
||||
cam_info["cam_to_world_R"],
|
||||
mask_R,
|
||||
display_table_mask_label,
|
||||
)["points_world"]
|
||||
display_pts_L = PtsUtil.random_downsample_point_cloud(
|
||||
display_pts_L, random_downsample_N
|
||||
)
|
||||
point_cloud_R = PtsUtil.random_downsample_point_cloud(
|
||||
display_pts_R, random_downsample_N
|
||||
)
|
||||
display_pts_overlap = DataLoadUtil.get_overlapping_points(
|
||||
display_pts_L, display_pts_R, voxel_size
|
||||
)
|
||||
return overlap_points, display_pts_overlap
|
||||
return overlap_points
|
||||
else:
|
||||
depth = DataLoadUtil.load_depth(
|
||||
@@ -390,27 +397,6 @@ class DataLoadUtil:
|
||||
)["points_world"]
|
||||
return point_cloud
|
||||
|
||||
@staticmethod
|
||||
def voxelize_points(points, voxel_size):
|
||||
|
||||
voxel_indices = np.floor(points / voxel_size).astype(np.int32)
|
||||
unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
|
||||
return unique_voxels
|
||||
|
||||
@staticmethod
|
||||
def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005):
|
||||
voxels_L, indices_L = DataLoadUtil.voxelize_points(point_cloud_L, voxel_size)
|
||||
voxels_R, _ = DataLoadUtil.voxelize_points(point_cloud_R, voxel_size)
|
||||
|
||||
voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
|
||||
voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
|
||||
overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
|
||||
mask_L = np.isin(
|
||||
indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
|
||||
)
|
||||
overlapping_points = point_cloud_L[mask_L]
|
||||
return overlapping_points
|
||||
|
||||
@staticmethod
|
||||
def load_points_normals(root, scene_name, display_table_as_world_space_origin=True):
|
||||
points_path = os.path.join(root, scene_name, "points_and_normals.txt")
|
||||
|
||||
Reference in New Issue
Block a user