This commit is contained in:
hofee 2024-11-02 21:47:01 +00:00
parent 9f0225c6d8
commit 8198515c7a

View File

@ -2,6 +2,7 @@ import os
import bpy import bpy
import sys import sys
import json import json
import time
import mathutils import mathutils
import numpy as np import numpy as np
sys.path.append(os.path.dirname(os.path.abspath(__file__))) sys.path.append(os.path.dirname(os.path.abspath(__file__)))
@ -14,8 +15,8 @@ class DataRenderer:
"renderer": { "renderer": {
"generate": "generate":
{ {
"object_dir": "/media/hofee/data/data/scaled_object_meshes", "object_dir": "/data/hofee/data/scaled_object_meshes",
"table_model_path": "/media/hofee/data/data/others/table.obj", "table_model_path": "/data/hofee/data/table.obj",
"plane_size": 10, "plane_size": 10,
"binocular_vision": True, "binocular_vision": True,
"light_and_camera_config":{ "light_and_camera_config":{
@ -42,7 +43,6 @@ class DataRenderer:
self.light_and_camera_config = config["renderer"]["generate"]["light_and_camera_config"] self.light_and_camera_config = config["renderer"]["generate"]["light_and_camera_config"]
self.obj_dir = config["renderer"]["generate"]["object_dir"] self.obj_dir = config["renderer"]["generate"]["object_dir"]
self.binocular_vision = config["renderer"]["generate"]["binocular_vision"] self.binocular_vision = config["renderer"]["generate"]["binocular_vision"]
self.obj_name_list = os.listdir(self.obj_dir)
self.target_obj = None self.target_obj = None
self.random_obj_list = [] self.random_obj_list = []
@ -58,15 +58,24 @@ class DataRenderer:
def do_render(self, cam_pose, restore_info, temp_dir): def do_render(self, cam_pose, restore_info, temp_dir):
self.reset() self.reset()
start_time = time.time()
self.restore_scene(restore_info=restore_info) self.restore_scene(restore_info=restore_info)
end_time = time.time()
print(f"Time taken for restoring scene: {end_time - start_time} seconds")
object_name = self.target_obj.name object_name = self.target_obj.name
temp_file_name = f"tmp" temp_file_name = f"tmp"
if "." in object_name: if "." in object_name:
object_name = object_name.split(".")[0] object_name = object_name.split(".")[0]
start_time = time.time()
BlenderUtils.set_camera_at(cam_pose) BlenderUtils.set_camera_at(cam_pose)
BlenderUtils.render_mask(temp_dir, temp_file_name, binocular_vision=self.binocular_vision) BlenderUtils.render_mask(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
BlenderUtils.render_normal_and_depth(temp_dir, temp_file_name, binocular_vision=self.binocular_vision) BlenderUtils.render_normal_and_depth(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
end_time = time.time()
print(f"Time taken for rendering: {end_time - start_time} seconds")
BlenderUtils.save_cam_params(temp_dir, temp_file_name, binocular_vision=self.binocular_vision) BlenderUtils.save_cam_params(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
depth_dir = os.path.join(temp_dir, "depth") depth_dir = os.path.join(temp_dir, "depth")
for depth_file in os.listdir(depth_dir): for depth_file in os.listdir(depth_dir):
if not depth_file.endswith(".png"): if not depth_file.endswith(".png"):
@ -112,9 +121,14 @@ def main(temp_dir):
scene_info_path = os.path.join(params_data["scene_path"], "scene_info.json") scene_info_path = os.path.join(params_data["scene_path"], "scene_info.json")
with open(scene_info_path, 'r') as f: with open(scene_info_path, 'r') as f:
scene_info = json.load(f) scene_info = json.load(f)
start_time = time.time()
data_renderer = DataRenderer() data_renderer = DataRenderer()
end_time = time.time()
print(f"Time taken for initialization: {end_time - start_time} seconds")
start_time = time.time()
data_renderer.do_render(cam_pose, scene_info, temp_dir) data_renderer.do_render(cam_pose, scene_info, temp_dir)
end_time = time.time()
print(f"Time taken for rendering: {end_time - start_time} seconds")
depth_dir = os.path.join(temp_dir, "depth") depth_dir = os.path.join(temp_dir, "depth")
for depth_file in os.listdir(depth_dir): for depth_file in os.listdir(depth_dir):
if not depth_file.endswith(".png"): if not depth_file.endswith(".png"):