finish task_gen
This commit is contained in:
363
task_gen_dependencies/solver_2d.py
Normal file
363
task_gen_dependencies/solver_2d.py
Normal file
@@ -0,0 +1,363 @@
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
from task_gen_dependencies.layout_2d import DFS_Solver_Floor
|
||||
from task_gen_dependencies.multi_add_util import *
|
||||
from task_gen_dependencies.utils import axis_to_quaternion, quaternion_rotate, get_rotation_matrix_from_quaternion, \
|
||||
get_quaternion_from_rotation_matrix, get_xyz_euler_from_quaternion
|
||||
from pyboot.utils.log import Log
|
||||
|
||||
def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True):
|
||||
"""
|
||||
根据指定的角度和旋转轴来旋转target_affine。
|
||||
参数:
|
||||
- target_affine: 4x4 仿射变换矩阵
|
||||
- angle_degrees: 旋转角度（以度为单位）
|
||||
- rot_axis: 旋转轴，'x'、'y' 或 'z'
|
||||
"""
|
||||
# 将角度转换为弧度
|
||||
angle_radians = np.deg2rad(angle_degrees)
|
||||
|
||||
# 创建旋转对象
|
||||
if rot_axis == 'z':
|
||||
rotation_vector = np.array([0, 0, angle_radians])
|
||||
elif rot_axis == 'y':
|
||||
rotation_vector = np.array([0, angle_radians, 0])
|
||||
elif rot_axis == 'x':
|
||||
rotation_vector = np.array([angle_radians, 0, 0])
|
||||
else:
|
||||
raise ValueError("Invalid rotation axis. Please choose from 'x', 'y', 'z'.")
|
||||
|
||||
# 生成旋转矩阵
|
||||
R_angle = R.from_rotvec(rotation_vector).as_matrix()
|
||||
|
||||
# 提取旋转部分（3x3矩阵）
|
||||
target_rotation = target_affine[:3, :3]
|
||||
|
||||
# 将 target_rotation 绕指定轴旋转指定角度，得到 target_rotation_2
|
||||
if use_local:
|
||||
target_rotation_2 = np.dot(target_rotation, R_angle)
|
||||
else:
|
||||
target_rotation_2 = np.dot(R_angle, target_rotation)
|
||||
|
||||
# 重新组合旋转矩阵 target_rotation_2 和原始的平移部分
|
||||
target_affine_2 = np.eye(4)
|
||||
target_affine_2[:3, :3] = target_rotation_2
|
||||
target_affine_2[:3, 3] = target_affine[:3, 3] # 保留原始的平移部分
|
||||
|
||||
return target_affine_2
|
||||
|
||||
|
||||
|
||||
|
||||
def quaternion_multiply(q1, q2):
|
||||
"""计算两个四元数的乘积"""
|
||||
w1, x1, y1, z1 = q1
|
||||
w2, x2, y2, z2 = q2
|
||||
|
||||
w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
|
||||
x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
|
||||
y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
|
||||
z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
|
||||
|
||||
return np.array([w, x, y, z])
|
||||
|
||||
def quaternion_rotate_z(quaternion, angle):
|
||||
"""
|
||||
Rotate a quaternion around the global z-axis by a given angle.
|
||||
|
||||
Parameters:
|
||||
quaternion (numpy array): The input quaternion [w, x, y, z].
|
||||
angle (float): The rotation angle in degrees.
|
||||
|
||||
Returns:
|
||||
numpy array: The rotated quaternion.
|
||||
"""
|
||||
# Convert angle from degrees to radians
|
||||
angle_rad = np.radians(angle)
|
||||
|
||||
# Calculate the rotation quaternion for z-axis rotation
|
||||
cos_half_angle = np.cos(angle_rad / 2)
|
||||
sin_half_angle = np.sin(angle_rad / 2)
|
||||
q_z = np.array([cos_half_angle, 0, 0, sin_half_angle])
|
||||
|
||||
# Rotate the input quaternion around the global z-axis
|
||||
rotated_quaternion = quaternion_multiply(q_z, quaternion)
|
||||
|
||||
return rotated_quaternion
|
||||
|
||||
|
||||
|
||||
def rotate_point_ext(px, py, angle, ox, oy):
|
||||
s, c = math.sin(angle), math.cos(angle)
|
||||
px, py = px - ox, py - oy
|
||||
xnew = px * c - py * s
|
||||
ynew = px * s + py * c
|
||||
return xnew + ox, ynew + oy
|
||||
|
||||
def get_corners(pose, size, angle):
|
||||
cx, cy = pose
|
||||
w, h = size
|
||||
corners = [
|
||||
rotate_point_ext(cx - w / 2, cy - h / 2, angle, cx, cy),
|
||||
rotate_point_ext(cx + w / 2, cy - h / 2, angle, cx, cy),
|
||||
rotate_point_ext(cx + w / 2, cy + h / 2, angle, cx, cy),
|
||||
rotate_point_ext(cx - w / 2, cy + h / 2, angle, cx, cy)
|
||||
]
|
||||
return corners
|
||||
|
||||
def compute_bounding_box(objects, expansion=40):
|
||||
all_corners = []
|
||||
for pose, size, angle in objects:
|
||||
all_corners.extend(get_corners(pose, size, angle))
|
||||
|
||||
min_x = min(x for x, y in all_corners)
|
||||
max_x = max(x for x, y in all_corners)
|
||||
min_y = min(y for x, y in all_corners)
|
||||
max_y = max(y for x, y in all_corners)
|
||||
|
||||
center_x = (min_x + max_x) / 2
|
||||
center_y = (min_y + max_y) / 2
|
||||
width = max_x - min_x + expansion
|
||||
height = max_y - min_y + expansion
|
||||
|
||||
return (center_x, center_y, width, height, 0)
|
||||
|
||||
|
||||
def compute_intersection(bbox, plane_center, plane_width, plane_height):
|
||||
# 解包最小外接矩形的信息
|
||||
bbox_center_x, bbox_center_y, bbox_width, bbox_height, _ = bbox
|
||||
|
||||
# 计算最小外接矩形的边界
|
||||
min_x = bbox_center_x - bbox_width / 2
|
||||
max_x = bbox_center_x + bbox_width / 2
|
||||
min_y = bbox_center_y - bbox_height / 2
|
||||
max_y = bbox_center_y + bbox_height / 2
|
||||
|
||||
# 计算平面的边界
|
||||
plane_min_x = plane_center[0] - plane_width / 2
|
||||
plane_max_x = plane_center[0] + plane_width / 2
|
||||
plane_min_y = plane_center[1] - plane_height / 2
|
||||
plane_max_y = plane_center[1] + plane_height / 2
|
||||
|
||||
# 计算相交部分的边界
|
||||
intersect_min_x = max(min_x, plane_min_x)
|
||||
intersect_max_x = min(max_x, plane_max_x)
|
||||
intersect_min_y = max(min_y, plane_min_y)
|
||||
intersect_max_y = min(max_y, plane_max_y)
|
||||
|
||||
# 检查相交区域是否有效
|
||||
if intersect_min_x < intersect_max_x and intersect_min_y < intersect_max_y:
|
||||
# 计算相交矩形的中心和尺寸
|
||||
intersect_center_x = (intersect_min_x + intersect_max_x) / 2
|
||||
intersect_center_y = (intersect_min_y + intersect_max_y) / 2
|
||||
intersect_width = intersect_max_x - intersect_min_x
|
||||
intersect_height = intersect_max_y - intersect_min_y
|
||||
|
||||
return (intersect_center_x, intersect_center_y, intersect_width, intersect_height,0)
|
||||
else:
|
||||
# 如果没有有效的相交区域，返回 None 或其他指示无相交的值
|
||||
return None
|
||||
|
||||
|
||||
class LayoutSolver2D:
|
||||
'''
|
||||
1. get room_vertices
|
||||
2. Generate Constraint with LLM
|
||||
3. Generate layout meet to the constraint
|
||||
'''
|
||||
|
||||
def __init__(self, workspace_xyz, workspace_size, objects=None, fix_obj_ids=[], obj_infos={}, angle_step=15):
|
||||
self.angle_step = angle_step
|
||||
x_half, y_half, z_half = workspace_size / 2
|
||||
room_vertices = [
|
||||
[-x_half, -y_half],
|
||||
[x_half, -y_half],
|
||||
[x_half, y_half],
|
||||
[-x_half, y_half]
|
||||
]
|
||||
self.plane_width = workspace_size[0]
|
||||
self.plane_height = workspace_size[1]
|
||||
self.room_vertices = room_vertices
|
||||
|
||||
self.objects = objects
|
||||
self.obj_infos = obj_infos
|
||||
|
||||
|
||||
self.cx, self.cy, self.cz = (coord * 1000 for coord in workspace_xyz)
|
||||
self.workspace_Z_half = workspace_size[2] / 2.0
|
||||
self.z_offset = 20
|
||||
self.fix_obj_ids = fix_obj_ids
|
||||
|
||||
def parse_solution(self, solutions, obj_id):
|
||||
[obj_cx, obj_cy], rotation, _ = solutions[obj_id][:3]
|
||||
obj_xyz = np.array([
|
||||
self.cx + obj_cx,
|
||||
self.cy + obj_cy,
|
||||
self.cz - self.workspace_Z_half + self.objects[obj_id].size[2]/2.0 + self.z_offset
|
||||
])
|
||||
init_quat = axis_to_quaternion(self.objects[obj_id].up_axis, "z")
|
||||
obj_quat = quaternion_rotate(init_quat, self.objects[obj_id].up_axis, -rotation)
|
||||
|
||||
obj_pose = np.eye(4)
|
||||
obj_pose[:3, :3] = get_rotation_matrix_from_quaternion(obj_quat)
|
||||
obj_pose[:3, 3] = obj_xyz
|
||||
self.objects[obj_id].obj_pose = obj_pose
|
||||
|
||||
|
||||
def old_solution(self,
|
||||
opt_obj_ids,
|
||||
exist_obj_ids,
|
||||
object_extent=50, # 外拓5cm
|
||||
start_with_edge=False,
|
||||
grid_size=0.01 # 1cm
|
||||
):
|
||||
solver = DFS_Solver_Floor(grid_size=int(grid_size*1000))
|
||||
room_poly = Polygon(self.room_vertices)
|
||||
grid_points = solver.create_grids(room_poly)
|
||||
|
||||
objs_succ = []
|
||||
saved_solutions = {} # TODO 将exist_obj_ids的pose保存进saved_solutions
|
||||
for obj_id in opt_obj_ids:
|
||||
size = self.objects[obj_id].size
|
||||
# import pdb;pdb.set_trace()
|
||||
size_extent = size[:2] + object_extent
|
||||
|
||||
|
||||
solutions = solver.get_all_solutions(room_poly, grid_points, size_extent)
|
||||
if len(solutions)>0:
|
||||
if start_with_edge:
|
||||
solutions = solver.place_edge(room_poly, solutions, size_extent)
|
||||
if len(saved_solutions) ==0:
|
||||
saved_solutions[obj_id] = random.choice(solutions)
|
||||
else:
|
||||
solutions = solver.filter_collision(saved_solutions, solutions)
|
||||
if len(solutions)>0:
|
||||
saved_solutions[obj_id] = random.choice(solutions)
|
||||
else:
|
||||
print(f"No valid solutions for apply {obj_id}.")
|
||||
continue
|
||||
|
||||
else:
|
||||
print(f"No valid solutions for apply {obj_id}.")
|
||||
continue
|
||||
|
||||
|
||||
self.parse_solution(saved_solutions, obj_id)
|
||||
objs_succ.append(obj_id)
|
||||
|
||||
return objs_succ
|
||||
|
||||
def __call__(self,
|
||||
opt_obj_ids,
|
||||
exist_obj_ids,
|
||||
object_extent=50, # 外拓5cm
|
||||
start_with_edge=False,
|
||||
key_obj = True,
|
||||
grid_size=0.01, # 1cm
|
||||
initial_angle = 0
|
||||
):
|
||||
|
||||
objs_succ = []
|
||||
fail_objs = []
|
||||
placed_objects = []
|
||||
main_objects = []
|
||||
label_flag = "no extra"
|
||||
|
||||
if not key_obj:
|
||||
for obj_id in self.objects:
|
||||
if obj_id not in opt_obj_ids:
|
||||
if obj_id in self.fix_obj_ids:
|
||||
continue
|
||||
size = self.objects[obj_id].size
|
||||
pose = self.objects[obj_id].obj_pose
|
||||
quaternion_rotate = get_quaternion_from_rotation_matrix(pose[:3, :3])
|
||||
angle_pa = get_xyz_euler_from_quaternion(quaternion_rotate)[2]
|
||||
main_objects.append(((pose[0, 3] - self.cx, pose[1, 3] - self.cy), (size[0], size[1]), angle_pa))
|
||||
# main_objects.append(((pose[0, 3] - self.cx, pose[1, 3] - self.cy), (size[0], size[1]), angle_pa * (180 / math.pi)))
|
||||
main_bounding_box_info = compute_bounding_box(main_objects, expansion=object_extent)
|
||||
|
||||
intersection = compute_intersection(main_bounding_box_info, (0, 0), self.plane_width, self.plane_height)
|
||||
if intersection is None:
|
||||
return objs_succ
|
||||
placed_objects.append((intersection, get_rotated_corners(*intersection)))
|
||||
label_flag = "add extra"
|
||||
object_extent = 0
|
||||
|
||||
obj_sizes = []
|
||||
grid_points = create_grid(self.plane_width, self.plane_height, int(grid_size*1000))
|
||||
saved_solutions = {}
|
||||
if len(opt_obj_ids) ==1:
|
||||
attempts = 800
|
||||
else:
|
||||
attempts = 400
|
||||
|
||||
for obj_id in opt_obj_ids:
|
||||
size = self.objects[obj_id].size
|
||||
_extent = self.obj_infos[obj_id].get("extent", object_extent)
|
||||
size_extent = size[:2] + _extent
|
||||
obj_sizes.append((size_extent[0],size_extent[1]))
|
||||
width, height = size_extent[0],size_extent[1]
|
||||
area_ratio = (width * height) / (self.plane_width * self.plane_height)
|
||||
|
||||
|
||||
# while len(placed_objects) < num_objects and attempts < max_attempts:
|
||||
# width, height = obj_sizes[len(placed_objects)]
|
||||
valid_position = False
|
||||
best_position = None
|
||||
max_distance = 1
|
||||
available_grid_points = filter_occupied_grids(grid_points, placed_objects)
|
||||
for idx in range(attempts):
|
||||
if not available_grid_points:
|
||||
break
|
||||
if len(opt_obj_ids) == 1 and area_ratio > 0.5:
|
||||
if idx < len(available_grid_points):
|
||||
x, y = available_grid_points[idx] # 使用索引获取点
|
||||
angle = random.choice([0,90,180,270])
|
||||
else:
|
||||
# 如果索引超出范围，可以选择退出循环或采取其他措施
|
||||
break
|
||||
else:
|
||||
x, y = random.choice(available_grid_points) # 随机选择一个点
|
||||
angle = np.random.choice(np.arange(0, 360, self.angle_step))
|
||||
# x, y = available_grid_points[idx] # random.choice(available_grid_points)
|
||||
# x = random.uniform(-self.plane_width / 2 + width / 2, self.plane_width / 2 - width / 2)
|
||||
# y = random.uniform(-self.plane_height / 2 + height / 2, self.plane_height / 2 - height / 2)
|
||||
# if len(placed_objects)==0:
|
||||
# angle = initial_angle
|
||||
# else:
|
||||
|
||||
|
||||
new_corners = get_rotated_corners(x, y, width, height, angle)
|
||||
|
||||
if is_within_bounds(new_corners, self.plane_width, self.plane_height) and not is_collision(new_corners, placed_objects):
|
||||
if not placed_objects:
|
||||
best_position = (x, y, width, height, angle)
|
||||
valid_position = True
|
||||
break
|
||||
|
||||
min_distance = min(calculate_distance(new_corners, obj[1]) for obj in placed_objects)
|
||||
if min_distance > max_distance:
|
||||
max_distance = min_distance
|
||||
best_position = (x, y, width, height, angle)
|
||||
valid_position = True
|
||||
# break
|
||||
|
||||
if valid_position:
|
||||
placed_objects.append((best_position, get_rotated_corners(*best_position)))
|
||||
saved_solutions[obj_id] = [[best_position[0],best_position[1]],best_position[4],None]
|
||||
self.parse_solution(saved_solutions, obj_id)
|
||||
objs_succ.append(obj_id)
|
||||
else:
|
||||
fail_objs.append(obj_id)
|
||||
|
||||
|
||||
# objects = [obj[0] for obj in placed_objects]
|
||||
# visualize_objects(objects, self.plane_width, self.plane_height,str(obj_id)+label_flag+".jpg")
|
||||
if len(fail_objs)>0:
|
||||
Log.warning("failed objects in layout 2d: " + str(fail_objs))
|
||||
|
||||
for obj_id in objs_succ:
|
||||
self.objects[obj_id].obj_pose = rotate_along_axis(self.objects[obj_id].obj_pose, random.uniform(-self.angle_step/2.0, self.angle_step/2.0), rot_axis='z', use_local=False)
|
||||
return objs_succ
|
||||
|
||||
|
||||
Reference in New Issue
Block a user