add min_cam_table_included_degree and random_view
This commit is contained in:
hofee
parent
70280a7b92
commit
f94cb2c0d6
@ -21,6 +21,8 @@ class DataGenerator:
|
||||
self.min_views = config["runner"]["generate"]["min_views"]
|
||||
self.min_diag = config["runner"]["generate"]["min_diag"]
|
||||
self.max_diag = config["runner"]["generate"]["max_diag"]
|
||||
self.min_cam_table_included_degree = config["runner"]["generate"]["min_cam_table_included_degree"]
|
||||
self.random_view_ratio = config["runner"]["generate"]["random_view_ratio"]
|
||||
self.binocular_vision = config["runner"]["generate"]["binocular_vision"]
|
||||
self.set_status_path = "http://localhost:5000/project/set_status"
|
||||
self.log_path = "http://localhost:5000/project/add_log"
|
||||
@ -240,7 +242,7 @@ class DataGenerator:
|
||||
if not os.path.exists(scene_dir):
|
||||
os.makedirs(scene_dir)
|
||||
view_num = int(self.min_views + (diag - self.min_diag)/(self.max_diag - self.min_diag) * (self.max_views - self.min_views))
|
||||
view_data = ViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.2,0.4), voxel_size=0.005, max_views=view_num)
|
||||
view_data = ViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.2,0.4), voxel_size=0.005, max_views=view_num, min_cam_table_included_degree = self.min_cam_table_included_degree, random_view_ratio = self.random_view_ratio )
|
||||
object_points = np.array(view_data["voxel_down_sampled_points"])
|
||||
normals = np.array(view_data["normals"])
|
||||
points_normals = np.concatenate((object_points, normals), axis=1)
|
||||
|
||||
151
pose.py
Normal file
151
pose.py
Normal file
@ -0,0 +1,151 @@
|
||||
import numpy as np
|
||||
|
||||
class PoseUtil:
|
||||
ROTATION = 1
|
||||
TRANSLATION = 2
|
||||
SCALE = 3
|
||||
|
||||
@staticmethod
|
||||
def get_uniform_translation(trans_m_min, trans_m_max, trans_unit, debug=False):
|
||||
if isinstance(trans_m_min, list):
|
||||
x_min, y_min, z_min = trans_m_min
|
||||
x_max, y_max, z_max = trans_m_max
|
||||
else:
|
||||
x_min, y_min, z_min = trans_m_min, trans_m_min, trans_m_min
|
||||
x_max, y_max, z_max = trans_m_max, trans_m_max, trans_m_max
|
||||
|
||||
x = np.random.uniform(x_min, x_max)
|
||||
y = np.random.uniform(y_min, y_max)
|
||||
z = np.random.uniform(z_min, z_max)
|
||||
translation = np.array([x, y, z])
|
||||
if trans_unit == "cm":
|
||||
translation = translation / 100
|
||||
if debug:
|
||||
print("uniform translation:", translation)
|
||||
return translation
|
||||
|
||||
@staticmethod
|
||||
def get_uniform_rotation(rot_degree_min=0, rot_degree_max=180, debug=False):
|
||||
axis = np.random.randn(3)
|
||||
axis /= np.linalg.norm(axis)
|
||||
theta = np.random.uniform(
|
||||
rot_degree_min / 180 * np.pi, rot_degree_max / 180 * np.pi
|
||||
)
|
||||
|
||||
K = np.array(
|
||||
[[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
|
||||
)
|
||||
R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * (K @ K)
|
||||
if debug:
|
||||
print("uniform rotation:", theta * 180 / np.pi)
|
||||
return R
|
||||
|
||||
@staticmethod
|
||||
def get_uniform_pose(
|
||||
trans_min, trans_max, rot_min=0, rot_max=180, trans_unit="cm", debug=False
|
||||
):
|
||||
translation = PoseUtil.get_uniform_translation(
|
||||
trans_min, trans_max, trans_unit, debug
|
||||
)
|
||||
rotation = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
||||
pose = np.eye(4)
|
||||
pose[:3, :3] = rotation
|
||||
pose[:3, 3] = translation
|
||||
return pose
|
||||
|
||||
@staticmethod
|
||||
def get_n_uniform_pose(
|
||||
trans_min,
|
||||
trans_max,
|
||||
rot_min=0,
|
||||
rot_max=180,
|
||||
n=1,
|
||||
trans_unit="cm",
|
||||
fix=None,
|
||||
contain_canonical=True,
|
||||
debug=False,
|
||||
):
|
||||
if fix == PoseUtil.ROTATION:
|
||||
translations = np.zeros((n, 3))
|
||||
for i in range(n):
|
||||
translations[i] = PoseUtil.get_uniform_translation(
|
||||
trans_min, trans_max, trans_unit, debug
|
||||
)
|
||||
if contain_canonical:
|
||||
translations[0] = np.zeros(3)
|
||||
rotations = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
||||
elif fix == PoseUtil.TRANSLATION:
|
||||
rotations = np.zeros((n, 3, 3))
|
||||
for i in range(n):
|
||||
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
||||
if contain_canonical:
|
||||
rotations[0] = np.eye(3)
|
||||
translations = PoseUtil.get_uniform_translation(
|
||||
trans_min, trans_max, trans_unit, debug
|
||||
)
|
||||
else:
|
||||
translations = np.zeros((n, 3))
|
||||
rotations = np.zeros((n, 3, 3))
|
||||
for i in range(n):
|
||||
translations[i] = PoseUtil.get_uniform_translation(
|
||||
trans_min, trans_max, trans_unit, debug
|
||||
)
|
||||
for i in range(n):
|
||||
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
||||
if contain_canonical:
|
||||
translations[0] = np.zeros(3)
|
||||
rotations[0] = np.eye(3)
|
||||
|
||||
pose = np.eye(4, 4, k=0)[np.newaxis, :].repeat(n, axis=0)
|
||||
pose[:, :3, :3] = rotations
|
||||
pose[:, :3, 3] = translations
|
||||
|
||||
return pose
|
||||
|
||||
@staticmethod
|
||||
def get_n_uniform_pose_batch(
|
||||
trans_min,
|
||||
trans_max,
|
||||
rot_min=0,
|
||||
rot_max=180,
|
||||
n=1,
|
||||
batch_size=1,
|
||||
trans_unit="cm",
|
||||
fix=None,
|
||||
contain_canonical=False,
|
||||
debug=False,
|
||||
):
|
||||
|
||||
batch_poses = []
|
||||
for i in range(batch_size):
|
||||
pose = PoseUtil.get_n_uniform_pose(
|
||||
trans_min,
|
||||
trans_max,
|
||||
rot_min,
|
||||
rot_max,
|
||||
n,
|
||||
trans_unit,
|
||||
fix,
|
||||
contain_canonical,
|
||||
debug,
|
||||
)
|
||||
batch_poses.append(pose)
|
||||
pose_batch = np.stack(batch_poses, axis=0)
|
||||
return pose_batch
|
||||
|
||||
@staticmethod
|
||||
def get_uniform_scale(scale_min, scale_max, debug=False):
|
||||
if isinstance(scale_min, list):
|
||||
x_min, y_min, z_min = scale_min
|
||||
x_max, y_max, z_max = scale_max
|
||||
else:
|
||||
x_min, y_min, z_min = scale_min, scale_min, scale_min
|
||||
x_max, y_max, z_max = scale_max, scale_max, scale_max
|
||||
|
||||
x = np.random.uniform(x_min, x_max)
|
||||
y = np.random.uniform(y_min, y_max)
|
||||
z = np.random.uniform(z_min, z_max)
|
||||
scale = np.array([x, y, z])
|
||||
if debug:
|
||||
print("uniform scale:", scale)
|
||||
return scale
|
||||
@ -3,6 +3,8 @@ import numpy as np
|
||||
import bmesh
|
||||
from collections import defaultdict
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
from blender.pose import PoseUtil
|
||||
import random
|
||||
|
||||
class ViewSampleUtil:
|
||||
@staticmethod
|
||||
@ -101,7 +103,7 @@ class ViewSampleUtil:
|
||||
return np.array(world_points), np.array(world_normals)
|
||||
|
||||
@staticmethod
|
||||
def get_cam_pose(view_data: dict, obj_world_pose: np.ndarray, max_views: int) -> np.ndarray:
|
||||
def get_cam_pose(view_data: dict, obj_world_pose: np.ndarray, max_views: int, min_cam_table_included_degree: int, random_view_ratio: float) -> np.ndarray:
|
||||
cam_poses = []
|
||||
min_height_z = 1000
|
||||
for look_at_point, cam_position in zip(view_data["look_at_points"], view_data["cam_positions"]):
|
||||
@ -134,7 +136,16 @@ class ViewSampleUtil:
|
||||
filtered_cam_poses = []
|
||||
for cam_pose in cam_poses:
|
||||
if cam_pose[2, 3] > min_height_z:
|
||||
direction_vector = cam_pose[:3, 2]
|
||||
horizontal_normal = np.array([0, 0, 1])
|
||||
cos_angle = np.dot(direction_vector, horizontal_normal) / (np.linalg.norm(direction_vector) * np.linalg.norm(horizontal_normal))
|
||||
angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))
|
||||
angle_degree = np.degrees(angle)
|
||||
if angle_degree < 90 - min_cam_table_included_degree:
|
||||
filtered_cam_poses.append(cam_pose)
|
||||
if random.random() < random_view_ratio:
|
||||
pertube_pose = PoseUtil.get_uniform_pose([0.1, 0.1, 0.1], [3, 3, 3], 0, 180, "cm")
|
||||
filtered_cam_poses.append(pertube_pose @ cam_pose)
|
||||
|
||||
if len(filtered_cam_poses) > max_views:
|
||||
indices = np.random.choice(len(filtered_cam_poses), max_views, replace=False)
|
||||
@ -143,10 +154,10 @@ class ViewSampleUtil:
|
||||
return np.array(filtered_cam_poses)
|
||||
|
||||
@staticmethod
|
||||
def sample_view_data_world_space(obj, distance_range:tuple = (0.3,0.5), voxel_size:float = 0.005, max_views: int=1) -> dict:
|
||||
def sample_view_data_world_space(obj, distance_range:tuple = (0.3,0.5), voxel_size:float = 0.005, max_views: int=1, min_cam_table_included_degree:int=20, random_view_ratio:float = 0.2) -> dict:
|
||||
obj_world_pose = np.asarray(obj.matrix_world)
|
||||
view_data = ViewSampleUtil.sample_view_data(obj, distance_range, voxel_size, max_views)
|
||||
view_data["cam_poses"] = ViewSampleUtil.get_cam_pose(view_data, obj_world_pose, max_views)
|
||||
view_data["cam_poses"] = ViewSampleUtil.get_cam_pose(view_data, obj_world_pose, max_views, min_cam_table_included_degree)
|
||||
view_data["voxel_down_sampled_points"], view_data["normals"] = ViewSampleUtil.get_world_points_and_normals(view_data, obj_world_pose)
|
||||
return view_data
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user