Compare commits
No commits in common. "c2849ce9bb3d83ac5084dd9a0624aa2eb53311ea" and "7d25777983af3da3373297471ee21c6b62591460" have entirely different histories.
c2849ce9bb
...
7d25777983
@ -1,86 +0,0 @@
|
|||||||
|
|
||||||
import os
|
|
||||||
import bpy
|
|
||||||
import numpy as np
|
|
||||||
from utils.blender_util import BlenderUtils
|
|
||||||
from utils.cad_view_sample_util import CADViewSampleUtil
|
|
||||||
from utils.material_util import MaterialUtil
|
|
||||||
|
|
||||||
class DataGenerator:
|
|
||||||
def __init__(self, config):
|
|
||||||
self.plane_size = config["runner"]["generate"]["plane_size"]
|
|
||||||
self.output_dir = config["runner"]["generate"]["output_dir"]
|
|
||||||
self.random_config = config["runner"]["generate"]["random_config"]
|
|
||||||
self.light_and_camera_config = config["runner"]["generate"]["light_and_camera_config"]
|
|
||||||
self.max_views = config["runner"]["generate"]["max_views"]
|
|
||||||
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.binocular_vision = config["runner"]["generate"]["binocular_vision"]
|
|
||||||
self.target_obj = None
|
|
||||||
self.stopped = False
|
|
||||||
self.random_obj_list = []
|
|
||||||
self.display_table_config = {}
|
|
||||||
BlenderUtils.setup_scene(self.light_and_camera_config, None, self.binocular_vision)
|
|
||||||
self.table = BlenderUtils.get_obj(BlenderUtils.TABLE_NAME)
|
|
||||||
|
|
||||||
def put_display_object(self, name):
|
|
||||||
obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, name)
|
|
||||||
obj = BlenderUtils.load_obj(name, obj_mesh_path)
|
|
||||||
bpy.ops.rigidbody.object_add()
|
|
||||||
bpy.context.object.rigid_body.type = 'ACTIVE'
|
|
||||||
MaterialUtil.change_object_material(obj, MaterialUtil.create_mask_material(color=(0, 1.0, 0)))
|
|
||||||
self.target_obj = obj
|
|
||||||
|
|
||||||
def reset(self):
|
|
||||||
self.target_obj = None
|
|
||||||
self.random_obj_list = []
|
|
||||||
BlenderUtils.reset_objects_and_platform()
|
|
||||||
|
|
||||||
def start_render(self, diag=0):
|
|
||||||
object_name = self.target_obj.name
|
|
||||||
if "." in object_name:
|
|
||||||
object_name = object_name.split(".")[0]
|
|
||||||
scene_dir = os.path.join(self.output_dir, object_name)
|
|
||||||
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 = CADViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.25,0.5), 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)
|
|
||||||
|
|
||||||
np.savetxt(os.path.join(scene_dir, "points_and_normals.txt"), points_normals)
|
|
||||||
for i, cam_pose in enumerate(view_data["cam_poses"]):
|
|
||||||
BlenderUtils.set_camera_at(cam_pose)
|
|
||||||
BlenderUtils.render_mask(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
|
|
||||||
BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision)
|
|
||||||
BlenderUtils.save_scene_info(scene_dir, self.display_table_config, object_name)
|
|
||||||
|
|
||||||
MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
|
|
||||||
|
|
||||||
for i, cam_pose in enumerate(view_data["cam_poses"]):
|
|
||||||
BlenderUtils.set_camera_at(cam_pose)
|
|
||||||
BlenderUtils.render_normal_and_depth(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
|
|
||||||
BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision)
|
|
||||||
|
|
||||||
depth_dir = os.path.join(scene_dir, "depth")
|
|
||||||
for depth_file in os.listdir(depth_dir):
|
|
||||||
if not depth_file.endswith(".png"):
|
|
||||||
name, _ = os.path.splitext(depth_file)
|
|
||||||
file_path = os.path.join(depth_dir, depth_file)
|
|
||||||
new_file_path = os.path.join(depth_dir, f"{name}.png")
|
|
||||||
os.rename(file_path,new_file_path)
|
|
||||||
BlenderUtils.save_blend(scene_dir)
|
|
||||||
exit(0)
|
|
||||||
|
|
||||||
return True
|
|
||||||
|
|
||||||
def gen_scene_data(self, object_name):
|
|
||||||
|
|
||||||
bpy.context.scene.frame_set(0)
|
|
||||||
self.put_display_object(object_name)
|
|
||||||
diag = BlenderUtils.get_obj_diag(self.target_obj.name)
|
|
||||||
self.start_render(diag)
|
|
||||||
|
|
||||||
|
|
@ -1,361 +0,0 @@
|
|||||||
import os
|
|
||||||
import json
|
|
||||||
import bpy
|
|
||||||
import time
|
|
||||||
import gc
|
|
||||||
import numpy as np
|
|
||||||
import mathutils
|
|
||||||
|
|
||||||
|
|
||||||
class CADBlenderUtils:
|
|
||||||
|
|
||||||
TABLE_NAME: str = "table"
|
|
||||||
CAMERA_NAME: str = "Camera"
|
|
||||||
CAMERA_RIGHT_NAME: str = "CameraRight"
|
|
||||||
CAMERA_OBJECT_NAME: str = "CameraObject"
|
|
||||||
DISPLAY_TABLE_NAME: str = "display_table"
|
|
||||||
MESH_FILE_NAME: str = "mesh.obj"
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_obj_path(obj_dir, name):
|
|
||||||
return os.path.join(obj_dir, name, CADBlenderUtils.MESH_FILE_NAME)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def load_obj(name, mesh_path, scale=1):
|
|
||||||
print(mesh_path)
|
|
||||||
bpy.ops.wm.obj_import(filepath=mesh_path)
|
|
||||||
loaded_object = bpy.context.selected_objects[-1]
|
|
||||||
loaded_object.name = name
|
|
||||||
loaded_object.data.name = name
|
|
||||||
loaded_object.scale = (scale, scale, scale)
|
|
||||||
bpy.ops.rigidbody.object_add()
|
|
||||||
return loaded_object
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_obj(name):
|
|
||||||
return bpy.data.objects.get(name)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_obj_pose(name):
|
|
||||||
obj = CADBlenderUtils.get_obj(name)
|
|
||||||
return np.asarray(obj.matrix_world)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def add_plane(name, location, orientation, size=10):
|
|
||||||
bpy.ops.mesh.primitive_plane_add(size=size, location=location)
|
|
||||||
plane = bpy.context.selected_objects[-1]
|
|
||||||
plane.name = name
|
|
||||||
plane.rotation_euler = orientation
|
|
||||||
bpy.ops.rigidbody.object_add()
|
|
||||||
bpy.context.object.rigid_body.type = "PASSIVE"
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def setup_scene(init_light_and_camera_config, table_model_path, binocular_vision):
|
|
||||||
bpy.context.scene.render.engine = "BLENDER_EEVEE"
|
|
||||||
bpy.context.scene.display.shading.show_xray = False
|
|
||||||
bpy.context.scene.display.shading.use_dof = False
|
|
||||||
bpy.context.scene.display.render_aa = "OFF"
|
|
||||||
bpy.context.scene.view_settings.view_transform = "Standard"
|
|
||||||
|
|
||||||
bpy.context.scene.eevee.use_ssr = False # 关闭屏幕空间反射
|
|
||||||
bpy.context.scene.eevee.use_bloom = False # 关闭辉光
|
|
||||||
bpy.context.scene.eevee.use_gtao = False # 关闭环境光遮蔽
|
|
||||||
bpy.context.scene.eevee.use_soft_shadows = False # 关闭软阴影
|
|
||||||
bpy.context.scene.eevee.use_shadows = False # 关闭所有阴影
|
|
||||||
bpy.context.scene.world.use_nodes = False # 如果你不需要环境光,关闭环境节点
|
|
||||||
|
|
||||||
# bpy.context.scene.eevee.use_sss = False # 关闭次表面散射
|
|
||||||
|
|
||||||
# 2. 设置最低的采样数
|
|
||||||
bpy.context.scene.eevee.taa_render_samples = 1
|
|
||||||
bpy.context.scene.eevee.taa_samples = 1
|
|
||||||
CADBlenderUtils.init_light_and_camera(
|
|
||||||
init_light_and_camera_config, binocular_vision
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def set_camera_params(camera, config, binocular_vision):
|
|
||||||
|
|
||||||
camera_object = bpy.data.objects.new(CADBlenderUtils.CAMERA_OBJECT_NAME, None)
|
|
||||||
bpy.context.collection.objects.link(camera_object)
|
|
||||||
cameras = [bpy.data.objects.get("Camera")]
|
|
||||||
camera.location = [0, 0, 0]
|
|
||||||
camera.rotation_euler = [0, 0, 0]
|
|
||||||
camera.parent = camera_object
|
|
||||||
if binocular_vision:
|
|
||||||
left_camera = cameras[0]
|
|
||||||
right_camera = left_camera.copy()
|
|
||||||
right_camera.name = CADBlenderUtils.CAMERA_RIGHT_NAME
|
|
||||||
right_camera.data = left_camera.data.copy()
|
|
||||||
right_camera.data.name = CADBlenderUtils.CAMERA_RIGHT_NAME
|
|
||||||
bpy.context.collection.objects.link(right_camera)
|
|
||||||
right_camera.parent = camera_object
|
|
||||||
right_camera.location = [config["eye_distance"] / 2, 0, 0]
|
|
||||||
left_camera.location = [-config["eye_distance"] / 2, 0, 0]
|
|
||||||
binocular_angle = config["eye_angle"]
|
|
||||||
half_angle = np.radians(binocular_angle / 2)
|
|
||||||
|
|
||||||
left_camera.rotation_euler[1] = -half_angle
|
|
||||||
right_camera.rotation_euler[1] = half_angle
|
|
||||||
cameras.append(right_camera)
|
|
||||||
|
|
||||||
for camera in cameras:
|
|
||||||
camera.data.clip_start = config["near_plane"]
|
|
||||||
camera.data.clip_end = config["far_plane"]
|
|
||||||
|
|
||||||
bpy.context.scene.render.resolution_x = config["resolution"][0]
|
|
||||||
bpy.context.scene.render.resolution_y = config["resolution"][1]
|
|
||||||
sensor_height = 24.0
|
|
||||||
focal_length = sensor_height / (
|
|
||||||
2 * np.tan(np.radians(config["fov_vertical"]) / 2)
|
|
||||||
)
|
|
||||||
camera.data.lens = focal_length
|
|
||||||
camera.data.sensor_width = (
|
|
||||||
sensor_height * config["resolution"][0] / config["resolution"][1]
|
|
||||||
)
|
|
||||||
camera.data.sensor_height = sensor_height
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def init_light_and_camera(init_light_and_camera_config, binocular_vision):
|
|
||||||
|
|
||||||
camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_NAME)
|
|
||||||
CADBlenderUtils.set_camera_params(
|
|
||||||
camera,
|
|
||||||
init_light_and_camera_config[CADBlenderUtils.CAMERA_NAME],
|
|
||||||
binocular_vision,
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_obj_diag(name):
|
|
||||||
obj = CADBlenderUtils.get_obj(name)
|
|
||||||
return np.linalg.norm(obj.dimensions)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def matrix_to_blender_pose(matrix):
|
|
||||||
location = matrix[:3, 3]
|
|
||||||
rotation_matrix = matrix[:3, :3]
|
|
||||||
rotation_matrix_blender = mathutils.Matrix(rotation_matrix.tolist())
|
|
||||||
rotation_euler = rotation_matrix_blender.to_euler()
|
|
||||||
return location, rotation_euler
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def set_camera_at(pose):
|
|
||||||
camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_OBJECT_NAME)
|
|
||||||
location, rotation_euler = CADBlenderUtils.matrix_to_blender_pose(pose)
|
|
||||||
|
|
||||||
camera.location = location
|
|
||||||
camera.rotation_euler = rotation_euler
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_object_bottom_z(obj):
|
|
||||||
vertices = [v.co for v in obj.data.vertices]
|
|
||||||
vertices_world = [obj.matrix_world @ v for v in vertices]
|
|
||||||
min_z = min([v.z for v in vertices_world])
|
|
||||||
return min_z
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def render_normal_and_depth(
|
|
||||||
output_dir, file_name, binocular_vision=False, target_object=None
|
|
||||||
):
|
|
||||||
# use pass z
|
|
||||||
bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
|
|
||||||
target_cameras = [CADBlenderUtils.CAMERA_NAME]
|
|
||||||
if binocular_vision:
|
|
||||||
target_cameras.append(CADBlenderUtils.CAMERA_RIGHT_NAME)
|
|
||||||
|
|
||||||
for cam_name in target_cameras:
|
|
||||||
bpy.context.scene.camera = CADBlenderUtils.get_obj(cam_name)
|
|
||||||
cam_suffix = "L" if cam_name == CADBlenderUtils.CAMERA_NAME else "R"
|
|
||||||
scene = bpy.context.scene
|
|
||||||
scene.render.filepath = ""
|
|
||||||
|
|
||||||
mask_dir = os.path.join(output_dir, "normal")
|
|
||||||
if not os.path.exists(mask_dir):
|
|
||||||
os.makedirs(mask_dir)
|
|
||||||
|
|
||||||
scene.render.filepath = os.path.join(
|
|
||||||
output_dir, mask_dir, f"{file_name}_{cam_suffix}.exr"
|
|
||||||
)
|
|
||||||
|
|
||||||
scene.render.image_settings.file_format = "OPEN_EXR"
|
|
||||||
scene.render.image_settings.color_mode = "RGB"
|
|
||||||
bpy.context.scene.view_settings.view_transform = "Raw"
|
|
||||||
scene.render.image_settings.color_depth = "16"
|
|
||||||
bpy.context.scene.render.filter_size = 1.5
|
|
||||||
scene.render.resolution_percentage = 100
|
|
||||||
scene.render.use_overwrite = False
|
|
||||||
scene.render.use_file_extension = False
|
|
||||||
scene.render.use_placeholder = False
|
|
||||||
scene.use_nodes = True
|
|
||||||
tree = scene.node_tree
|
|
||||||
|
|
||||||
for node in tree.nodes:
|
|
||||||
tree.nodes.remove(node)
|
|
||||||
|
|
||||||
rl = tree.nodes.new("CompositorNodeRLayers")
|
|
||||||
|
|
||||||
map_range = tree.nodes.new("CompositorNodeMapRange")
|
|
||||||
map_range.inputs["From Min"].default_value = 0.01
|
|
||||||
map_range.inputs["From Max"].default_value = 5
|
|
||||||
map_range.inputs["To Min"].default_value = 0
|
|
||||||
map_range.inputs["To Max"].default_value = 1
|
|
||||||
tree.links.new(rl.outputs["Depth"], map_range.inputs[0])
|
|
||||||
|
|
||||||
output_depth = tree.nodes.new("CompositorNodeOutputFile")
|
|
||||||
|
|
||||||
depth_dir = os.path.join(output_dir, "depth")
|
|
||||||
if not os.path.exists(depth_dir):
|
|
||||||
os.makedirs(depth_dir)
|
|
||||||
output_depth.base_path = depth_dir
|
|
||||||
output_depth.file_slots[0].path = f"{file_name}_{cam_suffix}.####"
|
|
||||||
output_depth.format.file_format = "PNG"
|
|
||||||
output_depth.format.color_mode = "BW"
|
|
||||||
output_depth.format.color_depth = "16"
|
|
||||||
tree.links.new(map_range.outputs[0], output_depth.inputs[0])
|
|
||||||
bpy.ops.render.render(write_still=True)
|
|
||||||
|
|
||||||
msg = "success"
|
|
||||||
return msg
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def render_mask(
|
|
||||||
output_dir, file_name, binocular_vision=False, target_object=None
|
|
||||||
):
|
|
||||||
target_cameras = [CADBlenderUtils.CAMERA_NAME]
|
|
||||||
if binocular_vision:
|
|
||||||
target_cameras.append(CADBlenderUtils.CAMERA_RIGHT_NAME)
|
|
||||||
|
|
||||||
for cam_name in target_cameras:
|
|
||||||
bpy.context.scene.camera = CADBlenderUtils.get_obj(cam_name)
|
|
||||||
cam_suffix = "L" if cam_name == CADBlenderUtils.CAMERA_NAME else "R"
|
|
||||||
scene = bpy.context.scene
|
|
||||||
scene.render.filepath = ""
|
|
||||||
|
|
||||||
mask_dir = os.path.join(output_dir, "mask")
|
|
||||||
if not os.path.exists(mask_dir):
|
|
||||||
os.makedirs(mask_dir)
|
|
||||||
|
|
||||||
scene.render.filepath = os.path.join(
|
|
||||||
output_dir, mask_dir, f"{file_name}_{cam_suffix}.png"
|
|
||||||
)
|
|
||||||
scene.render.image_settings.color_depth = "8"
|
|
||||||
scene.render.resolution_percentage = 100
|
|
||||||
scene.render.use_overwrite = False
|
|
||||||
scene.render.use_file_extension = False
|
|
||||||
scene.render.use_placeholder = False
|
|
||||||
|
|
||||||
|
|
||||||
bpy.ops.render.render(write_still=True)
|
|
||||||
|
|
||||||
msg = "success"
|
|
||||||
return msg
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def save_cam_params(scene_dir, idx, binocular_vision=False):
|
|
||||||
camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_NAME)
|
|
||||||
extrinsic = np.array(camera.matrix_world)
|
|
||||||
cam_data = camera.data
|
|
||||||
focal_length = cam_data.lens
|
|
||||||
sensor_width = cam_data.sensor_width
|
|
||||||
sensor_height = cam_data.sensor_height
|
|
||||||
resolution_x = bpy.context.scene.render.resolution_x
|
|
||||||
resolution_y = bpy.context.scene.render.resolution_y
|
|
||||||
intrinsic = np.zeros((3, 3))
|
|
||||||
intrinsic[0, 0] = focal_length * resolution_x / sensor_width # fx
|
|
||||||
intrinsic[1, 1] = focal_length * resolution_y / sensor_height # fy
|
|
||||||
intrinsic[0, 2] = resolution_x / 2.0 # cx
|
|
||||||
intrinsic[1, 2] = resolution_y / 2.0 # cy
|
|
||||||
intrinsic[2, 2] = 1.0
|
|
||||||
cam_object = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_OBJECT_NAME)
|
|
||||||
extrinsic_cam_object = np.array(cam_object.matrix_world)
|
|
||||||
data = {
|
|
||||||
"extrinsic": extrinsic.tolist(),
|
|
||||||
"extrinsic_cam_object": extrinsic_cam_object.tolist(),
|
|
||||||
"intrinsic": intrinsic.tolist(),
|
|
||||||
"far_plane": camera.data.clip_end,
|
|
||||||
"near_plane": camera.data.clip_start,
|
|
||||||
}
|
|
||||||
if binocular_vision:
|
|
||||||
right_camera = CADBlenderUtils.get_obj(CADBlenderUtils.CAMERA_RIGHT_NAME)
|
|
||||||
extrinsic_right = np.array(right_camera.matrix_world)
|
|
||||||
print("result:", extrinsic_right)
|
|
||||||
|
|
||||||
data["extrinsic_R"] = extrinsic_right.tolist()
|
|
||||||
|
|
||||||
cam_params_dir = os.path.join(scene_dir, "camera_params")
|
|
||||||
if not os.path.exists(cam_params_dir):
|
|
||||||
os.makedirs(cam_params_dir)
|
|
||||||
cam_params_path = os.path.join(cam_params_dir, f"{idx}.json")
|
|
||||||
with open(cam_params_path, "w") as f:
|
|
||||||
json.dump(data, f, indent=4)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def reset_objects_and_platform():
|
|
||||||
all_objects = bpy.data.objects
|
|
||||||
keep_objects = {
|
|
||||||
"plane_floor",
|
|
||||||
"plane_ceil",
|
|
||||||
"plane_wall_1",
|
|
||||||
"plane_wall_2",
|
|
||||||
"plane_wall_3",
|
|
||||||
"plane_wall_4",
|
|
||||||
}
|
|
||||||
keep_objects.add(CADBlenderUtils.CAMERA_OBJECT_NAME)
|
|
||||||
keep_objects.add(CADBlenderUtils.CAMERA_NAME)
|
|
||||||
keep_objects.add(CADBlenderUtils.CAMERA_RIGHT_NAME)
|
|
||||||
keep_objects.add(CADBlenderUtils.TABLE_NAME)
|
|
||||||
|
|
||||||
for obj in all_objects:
|
|
||||||
if obj.name not in keep_objects:
|
|
||||||
bpy.data.objects.remove(obj, do_unlink=True)
|
|
||||||
|
|
||||||
for block in bpy.data.meshes:
|
|
||||||
if block.users == 0:
|
|
||||||
bpy.data.meshes.remove(block)
|
|
||||||
for block in bpy.data.materials:
|
|
||||||
if block.users == 0:
|
|
||||||
bpy.data.materials.remove(block)
|
|
||||||
for block in bpy.data.images:
|
|
||||||
if block.users == 0:
|
|
||||||
bpy.data.images.remove(block)
|
|
||||||
|
|
||||||
gc.collect()
|
|
||||||
bpy.context.scene.frame_set(0)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def save_scene_info(scene_root_dir, display_table_config, target_name):
|
|
||||||
all_objects = bpy.data.objects
|
|
||||||
no_save_objects = {
|
|
||||||
"plane_floor",
|
|
||||||
"plane_ceil",
|
|
||||||
"plane_wall_1",
|
|
||||||
"plane_wall_2",
|
|
||||||
"plane_wall_3",
|
|
||||||
"plane_wall_4",
|
|
||||||
}
|
|
||||||
no_save_objects.add(CADBlenderUtils.CAMERA_OBJECT_NAME)
|
|
||||||
no_save_objects.add(CADBlenderUtils.CAMERA_NAME)
|
|
||||||
no_save_objects.add(CADBlenderUtils.CAMERA_RIGHT_NAME)
|
|
||||||
no_save_objects.add(CADBlenderUtils.TABLE_NAME)
|
|
||||||
scene_info = {}
|
|
||||||
for obj in all_objects:
|
|
||||||
if (
|
|
||||||
obj.name not in no_save_objects
|
|
||||||
and obj.name != CADBlenderUtils.DISPLAY_TABLE_NAME
|
|
||||||
):
|
|
||||||
obj_info = {
|
|
||||||
"location": list(obj.location),
|
|
||||||
"rotation_euler": list(obj.rotation_euler),
|
|
||||||
"scale": list(obj.scale),
|
|
||||||
}
|
|
||||||
scene_info[obj.name] = obj_info
|
|
||||||
scene_info[CADBlenderUtils.DISPLAY_TABLE_NAME] = display_table_config
|
|
||||||
scene_info["target_name"] = target_name
|
|
||||||
scene_info_path = os.path.join(scene_root_dir, "scene_info.json")
|
|
||||||
with open(scene_info_path, "w") as outfile:
|
|
||||||
json.dump(scene_info, outfile)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def save_blend(scene_root_dir):
|
|
||||||
blend_path = os.path.join(scene_root_dir, "scene.blend")
|
|
||||||
bpy.ops.wm.save_as_mainfile(filepath=blend_path)
|
|
@ -1,146 +0,0 @@
|
|||||||
|
|
||||||
import numpy as np
|
|
||||||
import bmesh
|
|
||||||
from collections import defaultdict
|
|
||||||
from scipy.spatial.transform import Rotation as R
|
|
||||||
from utils.pose import PoseUtil
|
|
||||||
from utils.pts import PtsUtil
|
|
||||||
import random
|
|
||||||
|
|
||||||
class CADViewSampleUtil:
|
|
||||||
@staticmethod
|
|
||||||
def farthest_point_sampling(points, num_samples):
|
|
||||||
num_points = points.shape[0]
|
|
||||||
if num_samples >= num_points:
|
|
||||||
return points, np.arange(num_points)
|
|
||||||
sampled_indices = np.zeros(num_samples, dtype=int)
|
|
||||||
sampled_indices[0] = np.random.randint(num_points)
|
|
||||||
min_distances = np.full(num_points, np.inf)
|
|
||||||
for i in range(1, num_samples):
|
|
||||||
current_point = points[sampled_indices[i - 1]]
|
|
||||||
dist_to_current_point = np.linalg.norm(points - current_point, axis=1)
|
|
||||||
min_distances = np.minimum(min_distances, dist_to_current_point)
|
|
||||||
sampled_indices[i] = np.argmax(min_distances)
|
|
||||||
downsampled_points = points[sampled_indices]
|
|
||||||
return downsampled_points, sampled_indices
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def voxel_downsample(points, voxel_size):
|
|
||||||
voxel_grid = defaultdict(list)
|
|
||||||
for i, point in enumerate(points):
|
|
||||||
voxel_index = tuple((point // voxel_size).astype(int))
|
|
||||||
voxel_grid[voxel_index].append(i)
|
|
||||||
|
|
||||||
downsampled_points = []
|
|
||||||
downsampled_indices = []
|
|
||||||
for indices in voxel_grid.values():
|
|
||||||
selected_index = indices[0]
|
|
||||||
downsampled_points.append(points[selected_index])
|
|
||||||
downsampled_indices.append(selected_index)
|
|
||||||
|
|
||||||
return np.array(downsampled_points), downsampled_indices
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def sample_view_data(obj, distance_range:tuple = (0.25,0.5), voxel_size:float = 0.005, max_views: int = 1, pertube_repeat:int = 1) -> dict:
|
|
||||||
view_data = {
|
|
||||||
"look_at_points": [],
|
|
||||||
"cam_positions": [],
|
|
||||||
}
|
|
||||||
mesh = obj.data
|
|
||||||
bm = bmesh.new()
|
|
||||||
bm.from_mesh(mesh)
|
|
||||||
bm.verts.ensure_lookup_table()
|
|
||||||
bm.faces.ensure_lookup_table()
|
|
||||||
bm.normal_update()
|
|
||||||
|
|
||||||
look_at_points = []
|
|
||||||
cam_positions = []
|
|
||||||
normals = []
|
|
||||||
for v in bm.verts:
|
|
||||||
look_at_point = np.array(v.co)
|
|
||||||
|
|
||||||
view_data["look_at_points"].append(look_at_point)
|
|
||||||
normal = np.zeros(3)
|
|
||||||
for loop in v.link_loops:
|
|
||||||
normal += np.array(loop.calc_normal())
|
|
||||||
normal /= len(v.link_loops)
|
|
||||||
normal = normal / np.linalg.norm(normal)
|
|
||||||
if np.isnan(normal).any():
|
|
||||||
continue
|
|
||||||
if np.dot(normal, look_at_point) < 0:
|
|
||||||
normal = -normal
|
|
||||||
normals.append(normal)
|
|
||||||
|
|
||||||
for _ in range(pertube_repeat):
|
|
||||||
perturb_angle = np.radians(np.random.uniform(0, 10))
|
|
||||||
perturb_axis = np.random.normal(size=3)
|
|
||||||
perturb_axis /= np.linalg.norm(perturb_axis)
|
|
||||||
rotation_matrix = R.from_rotvec(perturb_angle * perturb_axis).as_matrix()
|
|
||||||
perturbed_normal = np.dot(rotation_matrix, normal)
|
|
||||||
middle_distance = (distance_range[0] + distance_range[1]) / 2
|
|
||||||
perturbed_distance = random.uniform(middle_distance-0.05, middle_distance+0.05)
|
|
||||||
cam_position = look_at_point + perturbed_distance * perturbed_normal
|
|
||||||
look_at_points.append(look_at_point)
|
|
||||||
cam_positions.append(cam_position)
|
|
||||||
|
|
||||||
|
|
||||||
bm.free()
|
|
||||||
look_at_points = np.array(look_at_points)
|
|
||||||
cam_positions = np.array(cam_positions)
|
|
||||||
voxel_downsampled_look_at_points, selected_indices = CADViewSampleUtil.voxel_downsample(look_at_points, voxel_size)
|
|
||||||
voxel_downsampled_cam_positions = cam_positions[selected_indices]
|
|
||||||
voxel_downsampled_normals = np.array(normals)[selected_indices]
|
|
||||||
|
|
||||||
fps_downsampled_look_at_points, selected_indices = CADViewSampleUtil.farthest_point_sampling(voxel_downsampled_look_at_points, max_views*2)
|
|
||||||
fps_downsampled_cam_positions = voxel_downsampled_cam_positions[selected_indices]
|
|
||||||
|
|
||||||
view_data["look_at_points"] = fps_downsampled_look_at_points.tolist()
|
|
||||||
view_data["cam_positions"] = fps_downsampled_cam_positions.tolist()
|
|
||||||
view_data["normals"] = voxel_downsampled_normals
|
|
||||||
view_data["voxel_down_sampled_points"] = voxel_downsampled_look_at_points
|
|
||||||
return view_data
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def get_world_points_and_normals(view_data: dict, obj_world_pose: np.ndarray) -> tuple:
|
|
||||||
world_points = []
|
|
||||||
world_normals = []
|
|
||||||
for voxel_down_sampled_points, normal in zip(view_data["voxel_down_sampled_points"], view_data["normals"]):
|
|
||||||
voxel_down_sampled_points_world = obj_world_pose @ np.append(voxel_down_sampled_points, 1.0)
|
|
||||||
normal_world = obj_world_pose[:3, :3] @ normal
|
|
||||||
world_points.append(voxel_down_sampled_points_world[:3])
|
|
||||||
world_normals.append(normal_world)
|
|
||||||
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:
|
|
||||||
cam_poses = []
|
|
||||||
|
|
||||||
for look_at_point, cam_position in zip(view_data["look_at_points"], view_data["cam_positions"]):
|
|
||||||
look_at_point_world = obj_world_pose @ np.append(look_at_point, 1.0)
|
|
||||||
cam_position_world = obj_world_pose @ np.append(cam_position, 1.0)
|
|
||||||
look_at_point_world = look_at_point_world[:3]
|
|
||||||
cam_position_world = cam_position_world[:3]
|
|
||||||
forward_vector = cam_position_world - look_at_point_world
|
|
||||||
forward_vector /= np.linalg.norm(forward_vector)
|
|
||||||
up_vector = np.array([0, 0, 1])
|
|
||||||
right_vector = np.cross(up_vector, forward_vector)
|
|
||||||
rotation_matrix = np.array([right_vector, up_vector, forward_vector]).T
|
|
||||||
cam_pose = np.eye(4)
|
|
||||||
cam_pose[:3, :3] = rotation_matrix
|
|
||||||
cam_pose[:3, 3] = cam_position_world
|
|
||||||
cam_poses.append(cam_pose)
|
|
||||||
if len(cam_poses) > max_views:
|
|
||||||
cam_points = np.array([cam_pose[:3, 3] for cam_pose in cam_poses])
|
|
||||||
_, indices = PtsUtil.fps_downsample_point_cloud(cam_points, max_views, require_idx=True)
|
|
||||||
cam_poses = [cam_poses[i] for i in indices]
|
|
||||||
|
|
||||||
return np.array(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, min_cam_table_included_degree:int=20, random_view_ratio:float = 0.2) -> dict:
|
|
||||||
obj_world_pose = np.asarray(obj.matrix_world)
|
|
||||||
view_data = CADViewSampleUtil.sample_view_data(obj, distance_range, voxel_size, max_views)
|
|
||||||
view_data["cam_poses"] = CADViewSampleUtil.get_cam_pose(view_data, obj_world_pose, max_views, min_cam_table_included_degree, random_view_ratio)
|
|
||||||
view_data["voxel_down_sampled_points"], view_data["normals"] = CADViewSampleUtil.get_world_points_and_normals(view_data, obj_world_pose)
|
|
||||||
return view_data
|
|
||||||
|
|
@ -130,17 +130,14 @@ class ViewSampleUtil:
|
|||||||
|
|
||||||
dot_product = np.dot(forward_vector, up_vector)
|
dot_product = np.dot(forward_vector, up_vector)
|
||||||
angle = np.degrees(np.arccos(dot_product))
|
angle = np.degrees(np.arccos(dot_product))
|
||||||
right_vector = np.cross(up_vector, forward_vector)
|
|
||||||
print(angle)
|
print(angle)
|
||||||
if angle > 90 - min_cam_table_included_degree:
|
if angle < 110:
|
||||||
max_angle = 90 - min_cam_table_included_degree
|
target_angle = np.random.uniform(110, 150)
|
||||||
min_angle = max(90 - min_cam_table_included_degree*2, 30)
|
|
||||||
target_angle = np.random.uniform(min_angle, max_angle)
|
|
||||||
angle_difference = np.radians(target_angle - angle)
|
angle_difference = np.radians(target_angle - angle)
|
||||||
|
|
||||||
rotation_axis = np.cross(forward_vector, up_vector)
|
rotation_axis = np.cross(forward_vector, up_vector)
|
||||||
rotation_axis /= np.linalg.norm(rotation_axis)
|
rotation_axis /= np.linalg.norm(rotation_axis)
|
||||||
rotation_matrix = PoseUtil.rotation_matrix_from_axis_angle(rotation_axis, -angle_difference)
|
rotation_matrix = PoseUtil.rotation_matrix_from_axis_angle(rotation_axis, angle_difference)
|
||||||
new_cam_position_world = np.dot(rotation_matrix, cam_position_world - look_at_point_world) + look_at_point_world
|
new_cam_position_world = np.dot(rotation_matrix, cam_position_world - look_at_point_world) + look_at_point_world
|
||||||
cam_position_world = new_cam_position_world
|
cam_position_world = new_cam_position_world
|
||||||
forward_vector = cam_position_world - look_at_point_world
|
forward_vector = cam_position_world - look_at_point_world
|
||||||
@ -151,10 +148,7 @@ class ViewSampleUtil:
|
|||||||
corrected_up_vector = np.cross(forward_vector, right_vector)
|
corrected_up_vector = np.cross(forward_vector, right_vector)
|
||||||
rotation_matrix = np.array([right_vector, corrected_up_vector, forward_vector]).T
|
rotation_matrix = np.array([right_vector, corrected_up_vector, forward_vector]).T
|
||||||
else:
|
else:
|
||||||
right_vector = np.cross(up_vector, forward_vector)
|
rotation_matrix = np.array([right_vector, up_vector, forward_vector]).T
|
||||||
right_vector /= np.linalg.norm(right_vector)
|
|
||||||
corrected_up_vector = np.cross(forward_vector, right_vector)
|
|
||||||
rotation_matrix = np.array([right_vector, corrected_up_vector, forward_vector]).T
|
|
||||||
cam_pose = np.eye(4)
|
cam_pose = np.eye(4)
|
||||||
cam_pose[:3, :3] = rotation_matrix
|
cam_pose[:3, :3] = rotation_matrix
|
||||||
cam_pose[:3, 3] = cam_position_world
|
cam_pose[:3, 3] = cam_position_world
|
||||||
@ -168,7 +162,7 @@ class ViewSampleUtil:
|
|||||||
cos_angle = np.dot(direction_vector, horizontal_normal) / (np.linalg.norm(direction_vector) * np.linalg.norm(horizontal_normal))
|
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 = np.arccos(np.clip(cos_angle, -1.0, 1.0))
|
||||||
angle_degree = np.degrees(angle)
|
angle_degree = np.degrees(angle)
|
||||||
if angle_degree < 90 - min_cam_table_included_degree:
|
if angle_degree < 90 + min_cam_table_included_degree:
|
||||||
filtered_cam_poses.append(cam_pose)
|
filtered_cam_poses.append(cam_pose)
|
||||||
if random.random() < random_view_ratio:
|
if random.random() < random_view_ratio:
|
||||||
pertube_pose = PoseUtil.get_uniform_pose([0.1, 0.1, 0.1], [3, 3, 3], 0, 180, "cm")
|
pertube_pose = PoseUtil.get_uniform_pose([0.1, 0.1, 0.1], [3, 3, 3], 0, 180, "cm")
|
||||||
@ -178,7 +172,7 @@ class ViewSampleUtil:
|
|||||||
cam_points = np.array([cam_pose[:3, 3] for cam_pose in filtered_cam_poses])
|
cam_points = np.array([cam_pose[:3, 3] for cam_pose in filtered_cam_poses])
|
||||||
_, indices = PtsUtil.fps_downsample_point_cloud(cam_points, max_views, require_idx=True)
|
_, indices = PtsUtil.fps_downsample_point_cloud(cam_points, max_views, require_idx=True)
|
||||||
filtered_cam_poses = [filtered_cam_poses[i] for i in indices]
|
filtered_cam_poses = [filtered_cam_poses[i] for i in indices]
|
||||||
|
|
||||||
return np.array(filtered_cam_poses)
|
return np.array(filtered_cam_poses)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
Loading…
x
Reference in New Issue
Block a user