diff --git a/blender_util.py b/blender_util.py index 88d7676..9ce4faf 100644 --- a/blender_util.py +++ b/blender_util.py @@ -175,14 +175,39 @@ class BlenderUtils: @staticmethod def setup_render_normal(): bpy.context.scene.display.shading.light = 'MATCAP' + bpy.context.scene.display.shading.color_type = 'OBJECT' bpy.context.scene.display.shading.studio_light = 'check_normal+y.exr' @staticmethod - def render_and_save(output_dir, file_name, binocular_vision=False): + def render_and_save(output_dir, file_name, binocular_vision=False, target_object=None): target_cameras = [BlenderUtils.CAMERA_NAME] if binocular_vision: target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME) + + # BlenderUtils.setup_render_normal() + + # for cam_name in target_cameras: + # bpy.context.scene.camera = BlenderUtils.get_obj(cam_name) + # cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R" + # scene = bpy.context.scene + # scene.use_nodes = False + # scene.render.filepath = "" + # normal_dir = os.path.join(output_dir, "normal") + # if not os.path.exists(normal_dir): + # os.makedirs(normal_dir) + + + # scene.render.filepath = os.path.join(output_dir, normal_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) + + BlenderUtils.setup_render_mask() for cam_name in target_cameras: bpy.context.scene.camera = BlenderUtils.get_obj(cam_name) @@ -191,6 +216,7 @@ class BlenderUtils: 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) @@ -232,26 +258,7 @@ class BlenderUtils: bpy.ops.render.render(write_still=True) - BlenderUtils.setup_render_normal() - for cam_name in target_cameras: - bpy.context.scene.camera = BlenderUtils.get_obj(cam_name) - bpy.context.scene.view_layers["ViewLayer"].use_pass_normal = True - cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R" - scene = bpy.context.scene - scene.render.filepath = "" - - normal_dir = os.path.join(output_dir, "normal") - if not os.path.exists(normal_dir): - os.makedirs(normal_dir) - - scene.render.filepath = os.path.join(output_dir, normal_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 diff --git a/data_generator.py b/data_generator.py index d24d462..2019fc5 100644 --- a/data_generator.py +++ b/data_generator.py @@ -74,6 +74,7 @@ class DataGenerator: def add_log(self, msg, log_type): if not self.access: + print(f"[{log_type}] {msg}") return request_data = {"log":{}} request_data["log"]["message"] = msg @@ -156,47 +157,6 @@ class DataGenerator: obj.data.materials.append(mat) self.target_obj = obj - - def put_random_objects_on_table(self): - num_objects = self.random_config["random_objects"]["num"] - cluster = self.random_config["random_objects"]["cluster"] - for _ in range(num_objects): - obj_name = random.choice(self.obj_name_list) - print(obj_name) - obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, obj_name) - obj = BlenderUtils.load_obj(obj_name, obj_mesh_path) - - bbox = self.table.bound_box - bbox_world = [self.table.matrix_world @ mathutils.Vector(corner) for corner in bbox] - table_top_z = max([v.z for v in bbox_world]) - - platform_radius = self.platform.dimensions.x / 2.0 - - try_times = 0 - while True: - x = random.uniform(bbox_world[0].x*cluster, bbox_world[6].x*cluster) - y = random.uniform(bbox_world[0].y*cluster, bbox_world[6].y*cluster) - if math.sqrt(x**2 + y**2) > platform_radius*2 or try_times > 10: - break - try_times += 1 - if try_times > 10: - continue - - rotation = ( - random.uniform(0, 2 * np.pi), - random.uniform(0, 2 * np.pi), - random.uniform(0, 2 * np.pi) - ) - - obj_bottom_z = BlenderUtils.get_object_bottom_z(obj) - offset_z = obj_bottom_z - - obj.rotation_euler = rotation - obj.location = (x, y, table_top_z - offset_z) - - bpy.ops.rigidbody.object_add() - bpy.context.object.rigid_body.type = 'ACTIVE' - self.random_obj_list.append(obj) def reset(self): self.target_obj = None @@ -239,7 +199,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, min_cam_table_included_degree = self.min_cam_table_included_degree, random_view_ratio = self.random_view_ratio ) + view_data = ViewSampleUtil.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) @@ -247,7 +207,7 @@ class DataGenerator: 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_and_save(scene_dir, f"{i}", binocular_vision=self.binocular_vision) + BlenderUtils.render_and_save(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) self.set_progress("render frame", i, len(view_data["cam_poses"])) self.set_progress("render frame", len(view_data["cam_poses"]), len(view_data["cam_poses"])) @@ -304,8 +264,6 @@ class DataGenerator: if diag > self.max_diag or diag < self.min_diag: self.add_log(f"The diagonal size of the object <{object_name}>(size: {round(diag,3)}) does not meet the requirements.", "error") return "diag_error" - self.put_random_objects_on_table() - return self.simulate_scene(diag=diag) @@ -318,6 +276,12 @@ class DataGenerator: result = "retry" for target_obj_name in self.obj_name_list: self.add_log(f"Generating scene for object <{target_obj_name}>", "info") + scene_info_path = os.path.join(self.output_dir, target_obj_name, "scene_info.json") + if os.path.exists(scene_info_path): + self.add_log(f"Scene for object <{target_obj_name}> already exists, skipping", "warning") + count += 1 + continue + retry_times = 0 self.set_status("target_object", target_obj_name) while retry_times < 3 and result == "retry": @@ -341,4 +305,4 @@ class DataGenerator: self.set_progress("generate scene", count, total) result = "retry" - \ No newline at end of file + diff --git a/view_sample_util.py b/view_sample_util.py index 7d69e71..4e57f5c 100644 --- a/view_sample_util.py +++ b/view_sample_util.py @@ -40,7 +40,7 @@ class ViewSampleUtil: return np.array(downsampled_points), downsampled_indices @staticmethod - def sample_view_data(obj, distance_range:tuple = (0.2,0.4), voxel_size:float = 0.005, max_views: int = 1, pertube_repeat:int = 1) -> dict: + 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": [], @@ -76,9 +76,9 @@ class ViewSampleUtil: 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) - - distance = np.random.uniform(*distance_range) - cam_position = look_at_point + distance * perturbed_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)