fix scene_info

This commit is contained in:
hofee 2024-11-07 19:26:03 +08:00
parent 0ba46d4402
commit 378fd76f9f
4 changed files with 25 additions and 14 deletions

View File

@ -18,6 +18,9 @@ class DataGenerator:
self.random_config = config["runner"]["generate"]["random_config"]
self.light_and_camera_config = config["runner"]["generate"]["light_and_camera_config"]
self.obj_dir = config["runner"]["generate"]["object_dir"]
self.use_list = config["runner"]["generate"]["use_list"]
self.object_list_path = config["runner"]["generate"]["object_list_path"]
self.max_views = config["runner"]["generate"]["max_views"]
self.min_views = config["runner"]["generate"]["min_views"]
self.min_diag = config["runner"]["generate"]["min_diag"]
@ -30,7 +33,18 @@ class DataGenerator:
self.to_idx = config["runner"]["generate"]["to"]
self.set_status_path = f"http://localhost:{self.port}/project/set_status"
self.log_path = f"http://localhost:{self.port}/project/add_log"
self.obj_name_list = os.listdir(self.obj_dir)[self.from_idx: self.to_idx]
self.origin_obj_name_list = os.listdir(self.obj_dir)[self.from_idx: self.to_idx]
if not os.path.exists(self.output_dir):
os.makedirs(self.output_dir)
self.obj_name_list = []
if self.use_list:
self.target_name_list = [line.strip() for line in open(self.object_list_path).readlines()]
for obj_name in self.target_name_list:
if obj_name in self.origin_obj_name_list:
self.obj_name_list.append(obj_name)
else:
self.obj_name_list = self.origin_obj_name_list
self.target_obj = None
self.stopped = False
self.random_obj_list = []
@ -243,11 +257,7 @@ class DataGenerator:
msg = self.check_and_adjust_target()
if msg == "adjusted" and depth < 3:
bpy.context.view_layer.update()
bpy.context.scene.frame_set(0)
return self.simulate_scene(depth = depth + 1, diag=diag)
elif msg == "success":
if msg == "success":
print("Scene generation completed.")
result = self.start_render(diag=diag)
if not result:
@ -270,7 +280,7 @@ class DataGenerator:
def gen_all_scene_data(self):
max_retry_times = 3
max_retry_times = 5
total = len(self.obj_name_list)
count = 0
count_success = 0
@ -283,11 +293,12 @@ class DataGenerator:
if os.path.exists(scene_info_path):
self.add_log(f"Scene for object <{target_obj_name}> already exists, skipping", "warning")
count += 1
count_success += 1
continue
retry_times = 0
self.set_status("target_object", target_obj_name)
while retry_times < 3 and result == "retry":
while retry_times < max_retry_times and result == "retry":
self.reset()
try:
result = self.gen_scene_data(target_obj_name)

View File

@ -15,8 +15,8 @@ class DataRenderer:
"renderer": {
"generate":
{
"object_dir": "H:\\AI\\Datasets\\scaled_object_meshes",
"table_model_path": "H:\\AI\\Datasets\\table.obj",
"object_dir": "/media/hofee/data/data/scaled_object_meshes",
"table_model_path": "/media/hofee/data/data/others/table.obj",
"plane_size": 10,
"binocular_vision": True,
"light_and_camera_config":{
@ -144,4 +144,4 @@ if __name__ == "__main__":
print("Usage: blender -b -P data_renderer.py -- <temp_dir>")
else:
temp_dir = sys.argv[-1]
main(temp_dir)
main(temp_dir)

View File

@ -394,8 +394,8 @@ class BlenderUtils:
and obj.name != BlenderUtils.DISPLAY_TABLE_NAME
):
obj_info = {
"location": list(obj.location),
"rotation_euler": list(obj.rotation_euler),
"location": list(obj.matrix_world.translation),
"rotation_euler": list(obj.matrix_world.to_euler()),
"scale": list(obj.scale),
}
scene_info[obj.name] = obj_info

View File

@ -131,7 +131,7 @@ class ViewSampleUtil:
dot_product = np.dot(forward_vector, up_vector)
angle = np.degrees(np.arccos(dot_product))
right_vector = np.cross(up_vector, forward_vector)
print(angle)
if angle > 90 - min_cam_table_included_degree:
max_angle = 90 - min_cam_table_included_degree
min_angle = max(90 - min_cam_table_included_degree*2, 30)