add min_cam_table_included_degree and random_view

This commit is contained in:
hofee
2024-09-28 22:03:19 +08:00
parent 70280a7b92
commit f94cb2c0d6
3 changed files with 170 additions and 6 deletions

View File

@@ -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,8 +136,17 @@ class ViewSampleUtil:
filtered_cam_poses = []
for cam_pose in cam_poses:
if cam_pose[2, 3] > min_height_z:
filtered_cam_poses.append(cam_pose)
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)
filtered_cam_poses = [filtered_cam_poses[i] for i in indices]
@@ -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