137 lines
5.7 KiB
Python
137 lines
5.7 KiB
Python
import os
|
|
import bpy
|
|
import sys
|
|
import json
|
|
import mathutils
|
|
import numpy as np
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
from utils.blender_util import BlenderUtils
|
|
|
|
|
|
class DataRenderer:
|
|
def __init__(self):
|
|
config = {
|
|
"renderer": {
|
|
"generate":
|
|
{
|
|
"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":{
|
|
"Camera": {
|
|
"near_plane": 0.01,
|
|
"far_plane": 5,
|
|
"fov_vertical": 25,
|
|
"resolution": [1280, 800],
|
|
"eye_distance": 0.15,
|
|
"eye_angle": 25
|
|
},
|
|
"Light": {
|
|
"location": (0,0,3.5),
|
|
"orientation": (0, 0, 0),
|
|
"power": 150
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
self.plane_size = config["renderer"]["generate"]["plane_size"]
|
|
self.table_model_path = config["renderer"]["generate"]["table_model_path"]
|
|
self.light_and_camera_config = config["renderer"]["generate"]["light_and_camera_config"]
|
|
self.obj_dir = config["renderer"]["generate"]["object_dir"]
|
|
self.binocular_vision = config["renderer"]["generate"]["binocular_vision"]
|
|
self.obj_name_list = os.listdir(self.obj_dir)
|
|
self.target_obj = None
|
|
self.random_obj_list = []
|
|
|
|
BlenderUtils.setup_scene(self.light_and_camera_config, self.table_model_path, binocular_vision=self.binocular_vision)
|
|
self.table = BlenderUtils.get_obj(BlenderUtils.TABLE_NAME)
|
|
print(BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME))
|
|
|
|
def reset(self):
|
|
self.target_obj = None
|
|
self.random_obj_list = []
|
|
BlenderUtils.reset_objects_and_platform()
|
|
|
|
def do_render(self, cam_pose, restore_info, temp_dir):
|
|
|
|
self.reset()
|
|
self.restore_scene(restore_info=restore_info)
|
|
object_name = self.target_obj.name
|
|
temp_file_name = f"tmp"
|
|
if "." in object_name:
|
|
object_name = object_name.split(".")[0]
|
|
BlenderUtils.set_camera_at(cam_pose)
|
|
BlenderUtils.render_and_save(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
|
|
|
|
|
|
def restore_scene(self, restore_info):
|
|
|
|
for obj_name, obj_info in restore_info.items():
|
|
print("restoring", obj_name)
|
|
if obj_name == BlenderUtils.DISPLAY_TABLE_NAME:
|
|
self.restore_display_platform(obj_info)
|
|
else:
|
|
if obj_name == "target_name":
|
|
continue
|
|
if obj_name == restore_info["target_name"]:
|
|
obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, obj_name)
|
|
obj = BlenderUtils.load_obj(obj_name, obj_mesh_path)
|
|
obj.location = mathutils.Vector(obj_info["location"])
|
|
obj.rotation_euler = mathutils.Vector(obj_info["rotation_euler"])
|
|
obj.scale = mathutils.Vector(obj_info["scale"])
|
|
mat = bpy.data.materials.new(name="GreenMaterial")
|
|
mat.diffuse_color = (0.0, 1.0, 0.0, 1.0) # Green with full alpha (1.0)
|
|
if len(obj.data.materials) > 0:
|
|
obj.data.materials[0] = mat
|
|
else:
|
|
obj.data.materials.append(mat)
|
|
self.target_obj = obj
|
|
|
|
def restore_display_platform(self, platform_info):
|
|
bpy.ops.mesh.primitive_cylinder_add(radius=platform_info["radius"], depth=platform_info["height"])
|
|
platform = bpy.context.selected_objects[-1]
|
|
platform.name = BlenderUtils.DISPLAY_TABLE_NAME
|
|
platform.location = mathutils.Vector(platform_info["location"])
|
|
|
|
mat = bpy.data.materials.new(name="RedMaterial")
|
|
mat.diffuse_color = (1.0, 0.0, 0.0, 1.0) # Red with full alpha (1.0)
|
|
if len(platform.data.materials) > 0:
|
|
platform.data.materials[0] = mat
|
|
else:
|
|
platform.data.materials.append(mat)
|
|
|
|
bpy.ops.rigidbody.object_add()
|
|
bpy.context.object.rigid_body.type = 'PASSIVE'
|
|
bpy.ops.object.shade_auto_smooth()
|
|
|
|
def main(temp_dir):
|
|
params_data_path = os.path.join(temp_dir, "params.json")
|
|
with open(params_data_path, 'r') as f:
|
|
params_data = json.load(f)
|
|
cam_pose = np.array(params_data["cam_pose"])
|
|
print(cam_pose)
|
|
scene_info_path = os.path.join(params_data["scene_path"], "scene_info.json")
|
|
with open(scene_info_path, 'r') as f:
|
|
scene_info = json.load(f)
|
|
|
|
data_renderer = DataRenderer()
|
|
data_renderer.do_render(cam_pose, scene_info, temp_dir)
|
|
depth_dir = os.path.join(temp_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_cam_params(temp_dir, "tmp", binocular_vision=data_renderer.binocular_vision)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("Usage: blender -b -P data_renderer.py -- <temp_dir>")
|
|
else:
|
|
temp_dir = sys.argv[-1]
|
|
main(temp_dir) |