first commit
This commit is contained in:
391
utils/data_load.py
Normal file
391
utils/data_load.py
Normal file
@@ -0,0 +1,391 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import json
|
||||
import cv2
|
||||
import trimesh
|
||||
import torch
|
||||
import OpenEXR
|
||||
import Imath
|
||||
from utils.pts import PtsUtil
|
||||
|
||||
|
||||
class DataLoadUtil:
|
||||
TABLE_POSITION = np.asarray([0, 0, 0.8215])
|
||||
|
||||
@staticmethod
|
||||
def load_exr_image(file_path):
|
||||
exr_file = OpenEXR.InputFile(file_path)
|
||||
header = exr_file.header()
|
||||
dw = header['dataWindow']
|
||||
width = dw.max.x - dw.min.x + 1
|
||||
height = dw.max.y - dw.min.y + 1
|
||||
float_channels = ['R', 'G', 'B']
|
||||
img_data = []
|
||||
for channel in float_channels:
|
||||
channel_data = exr_file.channel(channel)
|
||||
img_data.append(np.frombuffer(channel_data, dtype=np.float16).reshape((height, width)))
|
||||
img = np.stack(img_data, axis=-1)
|
||||
return img
|
||||
|
||||
@staticmethod
|
||||
def get_display_table_info(root, scene_name):
|
||||
scene_info = DataLoadUtil.load_scene_info(root, scene_name)
|
||||
display_table_info = scene_info["display_table"]
|
||||
return display_table_info
|
||||
|
||||
@staticmethod
|
||||
def get_display_table_top(root, scene_name):
|
||||
display_table_height = DataLoadUtil.get_display_table_info(root, scene_name)[
|
||||
"height"
|
||||
]
|
||||
display_table_top = DataLoadUtil.TABLE_POSITION + np.asarray(
|
||||
[0, 0, display_table_height]
|
||||
)
|
||||
return display_table_top
|
||||
|
||||
@staticmethod
|
||||
def get_path(root, scene_name, frame_idx):
|
||||
path = os.path.join(root, scene_name, f"{frame_idx}")
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
def get_label_num(root, scene_name):
|
||||
label_dir = os.path.join(root, scene_name, "label")
|
||||
if not os.path.exists(label_dir):
|
||||
return 0
|
||||
return len(os.listdir(label_dir))
|
||||
|
||||
@staticmethod
|
||||
def get_label_path(root, scene_name, seq_idx):
|
||||
label_dir = os.path.join(root, scene_name, "label")
|
||||
if not os.path.exists(label_dir):
|
||||
os.makedirs(label_dir)
|
||||
path = os.path.join(label_dir, f"{seq_idx}.json")
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
def get_scene_seq_length(root, scene_name):
|
||||
camera_params_path = os.path.join(root, scene_name, "camera_params")
|
||||
return len(os.listdir(camera_params_path))
|
||||
|
||||
@staticmethod
|
||||
def load_mesh_at(model_dir, object_name, world_object_pose):
|
||||
model_path = os.path.join(model_dir, object_name, "mesh.obj")
|
||||
mesh = trimesh.load(model_path)
|
||||
mesh.apply_transform(world_object_pose)
|
||||
return mesh
|
||||
|
||||
@staticmethod
|
||||
def get_bbox_diag(model_dir, object_name):
|
||||
model_path = os.path.join(model_dir, object_name, "mesh.obj")
|
||||
mesh = trimesh.load(model_path)
|
||||
bbox = mesh.bounding_box.extents
|
||||
diagonal_length = np.linalg.norm(bbox)
|
||||
return diagonal_length
|
||||
|
||||
@staticmethod
|
||||
def load_scene_info(root, scene_name):
|
||||
scene_info_path = os.path.join(root, scene_name, "scene_info.json")
|
||||
with open(scene_info_path, "r") as f:
|
||||
scene_info = json.load(f)
|
||||
return scene_info
|
||||
|
||||
@staticmethod
|
||||
def load_target_pts_num_dict(root, scene_name):
|
||||
target_pts_num_path = os.path.join(root, scene_name, "target_pts_num.json")
|
||||
with open(target_pts_num_path, "r") as f:
|
||||
target_pts_num_dict = json.load(f)
|
||||
return target_pts_num_dict
|
||||
|
||||
@staticmethod
|
||||
def load_depth(path, min_depth=0.01, max_depth=5.0, binocular=False):
|
||||
|
||||
def load_depth_from_real_path(real_path, min_depth, max_depth):
|
||||
depth = cv2.imread(real_path, cv2.IMREAD_UNCHANGED)
|
||||
depth = depth.astype(np.float32) / 65535.0
|
||||
min_depth = min_depth
|
||||
max_depth = max_depth
|
||||
depth_meters = min_depth + (max_depth - min_depth) * depth
|
||||
return depth_meters
|
||||
|
||||
if binocular:
|
||||
depth_path_L = os.path.join(
|
||||
os.path.dirname(path), "depth", os.path.basename(path) + "_L.png"
|
||||
)
|
||||
depth_path_R = os.path.join(
|
||||
os.path.dirname(path), "depth", os.path.basename(path) + "_R.png"
|
||||
)
|
||||
depth_meters_L = load_depth_from_real_path(
|
||||
depth_path_L, min_depth, max_depth
|
||||
)
|
||||
depth_meters_R = load_depth_from_real_path(
|
||||
depth_path_R, min_depth, max_depth
|
||||
)
|
||||
return depth_meters_L, depth_meters_R
|
||||
else:
|
||||
depth_path = os.path.join(
|
||||
os.path.dirname(path), "depth", os.path.basename(path) + ".png"
|
||||
)
|
||||
depth_meters = load_depth_from_real_path(depth_path, min_depth, max_depth)
|
||||
return depth_meters
|
||||
|
||||
@staticmethod
|
||||
def load_seg(path, binocular=False, left_only=False):
|
||||
if binocular and not left_only:
|
||||
|
||||
def clean_mask(mask_image):
|
||||
green = [0, 255, 0]
|
||||
red = [255, 0, 0]
|
||||
threshold = 2
|
||||
mask_image = np.where(
|
||||
np.abs(mask_image - green) <= threshold, green, mask_image
|
||||
)
|
||||
mask_image = np.where(
|
||||
np.abs(mask_image - red) <= threshold, red, mask_image
|
||||
)
|
||||
return mask_image
|
||||
|
||||
mask_path_L = os.path.join(
|
||||
os.path.dirname(path), "mask", os.path.basename(path) + "_L.png"
|
||||
)
|
||||
mask_image_L = clean_mask(cv2.imread(mask_path_L, cv2.IMREAD_UNCHANGED))
|
||||
mask_path_R = os.path.join(
|
||||
os.path.dirname(path), "mask", os.path.basename(path) + "_R.png"
|
||||
)
|
||||
mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED))
|
||||
return mask_image_L, mask_image_R
|
||||
else:
|
||||
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_UNCHANGED)
|
||||
return mask_image
|
||||
|
||||
@staticmethod
|
||||
def load_normal(path, binocular=False, left_only=False, file_type="exr"):
|
||||
if binocular and not left_only:
|
||||
normal_path_L = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
|
||||
)
|
||||
normal_image_L = DataLoadUtil.load_exr_image(normal_path_L)
|
||||
|
||||
normal_path_R = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + f"_R.{file_type}"
|
||||
)
|
||||
normal_image_R = DataLoadUtil.load_exr_image(normal_path_R)
|
||||
normalized_normal_image_L = normal_image_L * 2.0 - 1.0
|
||||
normalized_normal_image_R = normal_image_R * 2.0 - 1.0
|
||||
return normalized_normal_image_L, normalized_normal_image_R
|
||||
else:
|
||||
if binocular and left_only:
|
||||
normal_path = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
|
||||
)
|
||||
else:
|
||||
normal_path = os.path.join(
|
||||
os.path.dirname(path), "normal", os.path.basename(path) + f".{file_type}"
|
||||
)
|
||||
normal_image = DataLoadUtil.load_exr_image(normal_path)
|
||||
normalized_normal_image = normal_image * 2.0 - 1.0
|
||||
return normalized_normal_image
|
||||
|
||||
@staticmethod
|
||||
def load_label(path):
|
||||
with open(path, "r") as f:
|
||||
label_data = json.load(f)
|
||||
return label_data
|
||||
|
||||
@staticmethod
|
||||
def load_from_preprocessed_pts(path, file_type="npy"):
|
||||
npy_path = os.path.join(
|
||||
os.path.dirname(path), "pts", os.path.basename(path) + "." + file_type
|
||||
)
|
||||
if file_type == "txt":
|
||||
pts = np.loadtxt(npy_path)
|
||||
else:
|
||||
pts = np.load(npy_path)
|
||||
return pts
|
||||
|
||||
@staticmethod
|
||||
def load_from_preprocessed_nrm(path, file_type="npy"):
|
||||
npy_path = os.path.join(
|
||||
os.path.dirname(path), "nrm", os.path.basename(path) + "." + file_type
|
||||
)
|
||||
if file_type == "txt":
|
||||
nrm = np.loadtxt(npy_path)
|
||||
else:
|
||||
nrm = np.load(npy_path)
|
||||
return nrm
|
||||
|
||||
@staticmethod
|
||||
def cam_pose_transformation(cam_pose_before):
|
||||
offset = np.asarray([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
|
||||
cam_pose_after = cam_pose_before @ offset
|
||||
return cam_pose_after
|
||||
|
||||
@staticmethod
|
||||
def load_cam_info(path, binocular=False, display_table_as_world_space_origin=True):
|
||||
scene_dir = os.path.dirname(path)
|
||||
root_dir = os.path.dirname(scene_dir)
|
||||
scene_name = os.path.basename(scene_dir)
|
||||
camera_params_path = os.path.join(
|
||||
os.path.dirname(path), "camera_params", os.path.basename(path) + ".json"
|
||||
)
|
||||
with open(camera_params_path, "r") as f:
|
||||
label_data = json.load(f)
|
||||
cam_to_world = np.asarray(label_data["extrinsic"])
|
||||
cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world)
|
||||
|
||||
if display_table_as_world_space_origin:
|
||||
world_to_display_table = np.eye(4)
|
||||
world_to_display_table[:3, 3] = -DataLoadUtil.get_display_table_top(
|
||||
root_dir, scene_name
|
||||
)
|
||||
cam_to_world = np.dot(world_to_display_table, cam_to_world)
|
||||
cam_intrinsic = np.asarray(label_data["intrinsic"])
|
||||
cam_info = {
|
||||
"cam_to_world": cam_to_world,
|
||||
"cam_intrinsic": cam_intrinsic,
|
||||
"far_plane": label_data["far_plane"],
|
||||
"near_plane": label_data["near_plane"],
|
||||
}
|
||||
if binocular:
|
||||
cam_to_world_R = np.asarray(label_data["extrinsic_R"])
|
||||
cam_to_world_R = DataLoadUtil.cam_pose_transformation(cam_to_world_R)
|
||||
cam_to_world_O = np.asarray(label_data["extrinsic_cam_object"])
|
||||
cam_to_world_O = DataLoadUtil.cam_pose_transformation(cam_to_world_O)
|
||||
if display_table_as_world_space_origin:
|
||||
cam_to_world_O = np.dot(world_to_display_table, cam_to_world_O)
|
||||
cam_to_world_R = np.dot(world_to_display_table, cam_to_world_R)
|
||||
cam_info["cam_to_world_O"] = cam_to_world_O
|
||||
cam_info["cam_to_world_R"] = cam_to_world_R
|
||||
return cam_info
|
||||
|
||||
@staticmethod
|
||||
def get_real_cam_O_from_cam_L(
|
||||
cam_L, cam_O_to_cam_L, scene_path, display_table_as_world_space_origin=True
|
||||
):
|
||||
root_dir = os.path.dirname(scene_path)
|
||||
scene_name = os.path.basename(scene_path)
|
||||
if isinstance(cam_L, torch.Tensor):
|
||||
cam_L = cam_L.cpu().numpy()
|
||||
nO_to_display_table_pose = cam_L @ cam_O_to_cam_L
|
||||
if display_table_as_world_space_origin:
|
||||
display_table_to_world = np.eye(4)
|
||||
display_table_to_world[:3, 3] = DataLoadUtil.get_display_table_top(
|
||||
root_dir, scene_name
|
||||
)
|
||||
nO_to_world_pose = np.dot(display_table_to_world, nO_to_display_table_pose)
|
||||
nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose)
|
||||
return nO_to_world_pose
|
||||
|
||||
@staticmethod
|
||||
def get_target_point_cloud(
|
||||
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")
|
||||
|
||||
z = depth
|
||||
x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
|
||||
y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
|
||||
|
||||
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
|
||||
mask = mask.reshape(-1, 4)
|
||||
|
||||
target_mask = (mask == target_mask_label).all(axis=-1)
|
||||
|
||||
target_points_camera = points_camera[target_mask]
|
||||
target_points_camera_aug = np.concatenate(
|
||||
[target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1
|
||||
)
|
||||
|
||||
target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
|
||||
data = {
|
||||
"points_world": target_points_world,
|
||||
"points_camera": target_points_camera,
|
||||
}
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def get_point_cloud(depth, cam_intrinsic, cam_extrinsic):
|
||||
h, w = depth.shape
|
||||
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
|
||||
|
||||
z = depth
|
||||
x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
|
||||
y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
|
||||
|
||||
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
|
||||
points_camera_aug = np.concatenate(
|
||||
[points_camera, np.ones((points_camera.shape[0], 1))], axis=-1
|
||||
)
|
||||
|
||||
points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
|
||||
return {"points_world": points_world, "points_camera": points_camera}
|
||||
|
||||
@staticmethod
|
||||
def get_target_point_cloud_world_from_path(
|
||||
path,
|
||||
binocular=False,
|
||||
random_downsample_N=65536,
|
||||
voxel_size=0.005,
|
||||
target_mask_label=(0, 255, 0, 255),
|
||||
display_table_mask_label=(0, 0, 255, 255),
|
||||
get_display_table_pts=False,
|
||||
require_normal=False,
|
||||
):
|
||||
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
|
||||
if binocular:
|
||||
depth_L, depth_R = DataLoadUtil.load_depth(
|
||||
path, cam_info["near_plane"], cam_info["far_plane"], binocular=True
|
||||
)
|
||||
mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
|
||||
point_cloud_L = DataLoadUtil.get_target_point_cloud(
|
||||
depth_L,
|
||||
cam_info["cam_intrinsic"],
|
||||
cam_info["cam_to_world"],
|
||||
mask_L,
|
||||
target_mask_label,
|
||||
)["points_world"]
|
||||
point_cloud_R = DataLoadUtil.get_target_point_cloud(
|
||||
depth_R,
|
||||
cam_info["cam_intrinsic"],
|
||||
cam_info["cam_to_world_R"],
|
||||
mask_R,
|
||||
target_mask_label,
|
||||
)["points_world"]
|
||||
point_cloud_L = PtsUtil.random_downsample_point_cloud(
|
||||
point_cloud_L, random_downsample_N
|
||||
)
|
||||
point_cloud_R = PtsUtil.random_downsample_point_cloud(
|
||||
point_cloud_R, random_downsample_N
|
||||
)
|
||||
overlap_points = PtsUtil.get_overlapping_points(
|
||||
point_cloud_L, point_cloud_R, voxel_size
|
||||
)
|
||||
return overlap_points
|
||||
else:
|
||||
depth = DataLoadUtil.load_depth(
|
||||
path, cam_info["near_plane"], cam_info["far_plane"]
|
||||
)
|
||||
mask = DataLoadUtil.load_seg(path)
|
||||
point_cloud = DataLoadUtil.get_target_point_cloud(
|
||||
depth, cam_info["cam_intrinsic"], cam_info["cam_to_world"], mask
|
||||
)["points_world"]
|
||||
return point_cloud
|
||||
|
||||
@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")
|
||||
points_normals = np.loadtxt(points_path)
|
||||
if display_table_as_world_space_origin:
|
||||
points_normals[:, :3] = points_normals[
|
||||
:, :3
|
||||
] - DataLoadUtil.get_display_table_top(root, scene_name)
|
||||
return points_normals
|
||||
Reference in New Issue
Block a user