scripts/json2urdf.py

import json
from xml.etree.ElementTree import Element, SubElement, tostring, ElementTree
from xml.dom.minidom import parseString
import pybullet as p
import pybullet_data
import os
from PIL import Image  # 使用Pillow保存图像
import numpy as np
import trimesh  # 用于处理3D网格
import imageio
import math

def degrees_to_radians(degrees):
    """Convert an angle from degrees to radians."""
    return degrees * math.pi / 180.0

def ply_to_obj(ply_filename, obj_filename, urdf_filename, part, json_data):
    """Convert a PLY file to OBJ format using trimesh."""
    base_path = '/'.join(urdf_filename.split('/')[:-1])
    mesh = trimesh.load(os.path.join(base_path, ply_filename), force='mesh')
    print(ply_filename, mesh.bounding_box.centroid)
    # mesh.vertices -= (mesh.bounding_box.centroid)
    # find base (parent == -1)
    base_part_id = next((p['id'] for p in json_data['diffuse_tree'] if p['parent'] == -1), None)
    if 'joint' in part.keys():
        mesh.vertices -= (mesh.bounding_box.centroid + part['joint']['axis']['origin'])
        while part['parent'] != base_part_id:
            parent_part = next((p for p in json_data['diffuse_tree'] if p['id'] == part['parent']), None)
            if parent_part is None:
                break
            mesh.vertices -= (parent_part['joint']['axis']['origin'])
            part = parent_part
    else:
        mesh.vertices -= (mesh.bounding_box.centroid)
    mesh.export(os.path.join(base_path, obj_filename))

def create_urdf_from_json(json_data, urdf_filename, parent_part=None):
    robot = Element('robot', name='articulate_object')

    def add_link(parent, part, urdf_filename, json_data, base=False):
        link = SubElement(parent, 'link', name=f"link_{part['id']}")
        
        # 将PLY文件转换为OBJ文件
        ply_path = part['objs'][0]
        obj_path = os.path.splitext(ply_path)[0] + '.obj'
        # if not os.path.exists(obj_path):
        ply_to_obj(ply_path, obj_path, urdf_filename, part, json_data)
        
        visual = SubElement(link, 'visual')
        origin = SubElement(visual, 'origin', xyz=" ".join(map(str, part['aabb']['center'])), rpy="0 0 0")
        geometry = SubElement(visual, 'geometry')
        mesh = SubElement(geometry, 'mesh', filename=obj_path)  # 使用转换后的OBJ路径
        

    def add_joint(parent, child_part, parent_part):
        joint = SubElement(parent, 'joint', name=f"{parent_part['id']}_{child_part['id']}_joint", type=child_part['joint']['type'])

        # for i in range(3):
        #     child_part['joint']['axis']['origin'][i] -= (child_part['aabb']['size'][i])
        origin = SubElement(joint, 'origin', xyz=" ".join(map(str, child_part['joint']['axis']['origin'])), rpy="0 0 0")
        # origin = SubElement(joint, 'origin', xyz=" 0 0 0 ", rpy="0 0 0")
        axis = SubElement(joint, 'axis', xyz=" ".join(map(str, child_part['joint']['axis']['direction'])))
        if child_part['joint']['type'] == 'revolute':
            child_part['joint']['range'][0] = degrees_to_radians(child_part['joint']['range'][0])
            child_part['joint']['range'][1] = degrees_to_radians(child_part['joint']['range'][1])
        
        lower, upper = child_part['joint']['range']
        if upper < lower:
            lower, upper = upper, lower

        limit = SubElement(
            joint, 'limit',
            lower=str(lower),
            upper=str(upper),
            effort="10",
            velocity="1"
        )
        parent_element = SubElement(joint, 'parent', link=f"link_{parent_part['id']}")
        child_element = SubElement(joint, 'child', link=f"link_{child_part['id']}")

    base_part = json_data['diffuse_tree'][0]
    add_link(robot, base_part, urdf_filename, json_data, base=True)

    for part in json_data['diffuse_tree'][1:]:
        base_part = next((p for p in json_data['diffuse_tree'] if p['parent'] == -1), None)
        parent_part = next((p for p in json_data['diffuse_tree'] if p['id'] == part['parent']), None)
        add_link(robot, part, urdf_filename, json_data)
        if parent_part:
            add_joint(robot, part, parent_part)

    xmlstr = parseString(tostring(robot)).toprettyxml(indent="   ")
    
    with open(urdf_filename, "w") as f:
        f.write(xmlstr)


def pybullet_render(urdf_path, target_dir, num_frames, distance=3, fov=60):
    physicsClient = p.connect(p.DIRECT)
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    try:
        robot = p.loadURDF(urdf_path, [0, 0, 0])
    except Exception as e:
        print(e)
        return
    # Disable part mesh colors.
    # for i in range(-1, p.getNumJoints(robot)):
    #     rgba = [np.random.uniform(0.2, 1.0), np.random.uniform(0.2, 1.0), np.random.uniform(0.2, 1.0), 1]
    #     p.changeVisualShape(robot, linkIndex=i, rgbaColor=rgba)
    p.resetBasePositionAndOrientation(robot, [0, 0, 0], [0, 0.7071, 0.7071, 0])
    joint_info = []
    for i in range(p.getNumJoints(robot)):
        info = p.getJointInfo(robot, i)
        if info[2] != p.JOINT_FIXED:
            joint_info.append({
                'index': info[0],
                'type': info[2],
                'name': info[1].decode('utf-8'),
                'lower_limit': info[8],
                'upper_limit': info[9],
                'initial_position': p.getJointState(robot, info[0])[0]
            })

    joint_positions = {}
    for joint in joint_info:
        start = joint['lower_limit']
        end = joint['upper_limit']
        joint_positions[joint['index']] = np.concatenate((np.linspace(start, end, num_frames), np.linspace(end, start, num_frames)))

    gif_frames = []
    for frame in range(num_frames*2):
        # import pdb; pdb.set_trace()
        for joint in joint_info:
            p.resetJointState(robot, joint['index'], joint_positions[joint['index']][frame])
            joint_state = p.getJointState(robot, joint['index'])
            p.stepSimulation()
        viewMatrix = p.computeViewMatrixFromYawPitchRoll(
            cameraTargetPosition=[0, 0, 0],
            distance=3.0,
            yaw=-150,
            pitch=-10,
            roll=0,
            upAxisIndex=2
        )
        projectionMatrix=p.computeProjectionMatrixFOV(
            fov=fov, ##60
            aspect=1.0, 
            nearVal=0.1, farVal=100)

        width, height, rgbPixels, depthBuffer, segMask = p.getCameraImage(
            width=1024, height=1024, viewMatrix=viewMatrix,
            projectionMatrix=projectionMatrix,
            renderer=p.ER_BULLET_HARDWARE_OPENGL)
    
        #get rgba image 
        rgba_image = np.reshape(rgbPixels, (height, width, 4))
        # rgba_image[np.all(rgba_image[:, :, :3] == 255, axis=-1)] = [0, 0, 0, 0]
        gif_frames.append(rgba_image[:, :, :3])
        
    p.disconnect()
    imageio.mimsave(f'{target_dir}/animation.gif', gif_frames, fps=8, loop=0)