87 lines
4.3 KiB
Python
87 lines
4.3 KiB
Python
|
|
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)
|
|
|
|
|