CoolFace
Apppublic

mj125/urdf-builder

sourceHugging Faceupdated 7mo agoView on Hugging Face
0likes
robot_renderer.py468 linesDownload Raw Back to root
1import numpy as np2import xml.etree.ElementTree as ET3import json4import html5import base646import os7 8class RobotRenderer:9    def __init__(self):10        self.joints = {}11        self.links = {}12        self.tree = {}13        self.base_link = None14        self.ordered_joints = []     # 움직이는 조인트 (슬라이더용)15        self.all_joint_names = []    # 모든 조인트 (체크박스용)16 17    def _parse_geometry(self, element, mesh_map, is_collision=False):18        info = {'type': None, 'dim': [0.1, 0.1, 0.1], 'scale': [1, 1, 1]}19        20        geom = element.find('geometry')21        if geom is None: return None22 23        if geom.find('mesh') is not None:24            if is_collision: return None 25            26            mesh_tag = geom.find('mesh')27            filename = mesh_tag.get('filename', '')28            29            # [Smart Matching] 대소문자 무시30            target_filename = os.path.basename(filename).lower()31            32            if target_filename in mesh_map:33                real_path = mesh_map[target_filename]34                try:35                    with open(real_path, "rb") as f:36                        ext = os.path.splitext(target_filename)[1]37                        if ext == '.obj':38                            info['type'] = 'obj'39                            info['mesh_data'] = f.read().decode('utf-8', errors='ignore')40                        else:41                            # STL (Binary)42                            info['type'] = 'mesh'43                            info['mesh_data'] = base64.b64encode(f.read()).decode('utf-8')44 45                        scale = mesh_tag.get('scale', '1 1 1').split()46                        info['scale'] = [float(s) for s in scale]47                        print(f"✅ Loaded Mesh: {filename}")48                except Exception as e:49                    print(f"❌ Read Error {filename}: {e}")50                    info['type'] = 'error_box'51            else:52                print(f"⚠️ Missing Mesh: {filename}")53                info['type'] = 'error_box'54 55        elif geom.find('cylinder') is not None:56            cyl = geom.find('cylinder')57            info['type'] = 'cylinder'58            info['dim'] = [float(cyl.get('radius', 0.05)), float(cyl.get('length', 0.1))]59            60        elif geom.find('box') is not None:61            box = geom.find('box')62            info['type'] = 'box'63            info['dim'] = [float(x) for x in box.get('size', '0.1 0.1 0.1').split()]64            65        elif geom.find('sphere') is not None:66            sph = geom.find('sphere')67            info['type'] = 'sphere'68            info['dim'] = [float(sph.get('radius', 0.05))]69            70        return info71 72    def load_urdf(self, urdf_content, mesh_files=[]):73        print("\n--- Loading URDF (V7 All Joints) ---")74        self.joints = {}75        self.links = {}76        self.tree = {}77        self.base_link = None78        self.ordered_joints = []79        self.all_joint_names = [] 80 81        mesh_map = {}82        if mesh_files:83            for f_path in mesh_files:84                fname = os.path.basename(f_path).lower()85                mesh_map[fname] = f_path86            print(f"📂 Uploaded Files Map: {list(mesh_map.keys())}")87 88        try: root = ET.fromstring(urdf_content)89        except ET.ParseError: return False90 91        # Link Parsing92        for i, link in enumerate(root.findall('link')):93            name = link.get('name')94            95            # 홀짝 기본 색상 (urdf_generator와 싱크를 맞춤)96            color_val = 0.8 if i % 2 == 0 else 0.497            default_color = [color_val, color_val, color_val]98            99            link_info = {'name': name, 'visual': None, 'collision': None, 'color': default_color}100            101            visual = link.find('visual')102            if visual is not None:103                vis_data = self._parse_geometry(visual, mesh_map, is_collision=False)104                if vis_data:105                    origin = visual.find('origin')106                    if origin is not None:107                        xyz = [float(x) for x in origin.get('xyz', '0 0 0').split()]108                        rpy = [float(x) for x in origin.get('rpy', '0 0 0').split()]109                        vis_data['origin'] = xyz + rpy110                    else: vis_data['origin'] = [0,0,0,0,0,0]111                    112                    material = visual.find('material')113                    if material is not None:114                        color = material.find('color')115                        if color is not None:116                            rgba = [float(x) for x in color.get('rgba', '0.5 0.5 0.5 1').split()]117                            link_info['color'] = rgba[:3]118                    vis_data['color'] = link_info['color']119                    link_info['visual'] = vis_data120 121            collision = link.find('collision')122            if collision is not None:123                col_data = self._parse_geometry(collision, mesh_map, is_collision=True)124                if col_data:125                    origin = collision.find('origin')126                    if origin is not None:127                        xyz = [float(x) for x in origin.get('xyz', '0 0 0').split()]128                        rpy = [float(x) for x in origin.get('rpy', '0 0 0').split()]129                        col_data['origin'] = xyz + rpy130                    else: col_data['origin'] = [0,0,0,0,0,0]131                    link_info['collision'] = col_data132 133            self.links[name] = link_info134 135        # Joint Parsing136        for joint in root.findall('joint'):137            name = joint.get('name'); type_ = joint.get('type')138            parent = joint.find('parent').get('link'); child = joint.find('child').get('link')139            140            origin = joint.find('origin')141            xyz = [float(x) for x in origin.get('xyz', '0 0 0').split()] if origin is not None else [0,0,0]142            rpy = [float(x) for x in origin.get('rpy', '0 0 0').split()] if origin is not None else [0,0,0]143            144            axis_elem = joint.find('axis')145            axis = [float(x) for x in axis_elem.get('xyz', '1 0 0').split()] if axis_elem is not None else [1,0,0]146            147            limit = joint.find('limit')148            lower, upper = -3.14, 3.14149            if limit is not None:150                if limit.get('lower'): lower = float(limit.get('lower'))151                if limit.get('upper'): upper = float(limit.get('upper'))152            153            self.joints[name] = {'parent': parent, 'child': child, 'xyz': xyz, 'rpy': rpy, 'axis': axis, 'limits': [np.degrees(lower), np.degrees(upper)], 'type': type_}154            155            self.all_joint_names.append(name)156            157            if type_ != 'fixed': self.ordered_joints.append(name)158            if parent not in self.tree: self.tree[parent] = []159            self.tree[parent].append(name)160 161        children = set(j['child'] for j in self.joints.values())162        roots = list(set(self.links.keys()) - children)163        if roots: self.base_link = roots[0]164        elif 'world' in self.links: self.base_link = 'world'165        166        print(f"✅ URDF Loaded. All Joints: {len(self.all_joint_names)} (Moving: {len(self.ordered_joints)})")167        return True168 169    def get_joint_list(self):170        res = []171        for name in self.all_joint_names:172            j = self.joints[name]173            res.append({174                'name': name,175                'type': j['type'],176                'parent': j['parent'],177                'child': j['child'],178                'min': j['limits'][0],179                'max': j['limits'][1],180                'axis': j['axis']181            })182        return res183 184    def get_viewer_html(self):185        robot_data = {'links': self.links, 'joints': self.joints, 'tree': self.tree, 'base': self.base_link, 'joint_order': self.ordered_joints}186        json_data = json.dumps(robot_data)187        188        html_code = f"""189        <!DOCTYPE html>190        <html><head>191        <style>192            body {{ margin: 0; overflow: hidden; font-family: monospace; }}193            #ui {{ position: absolute; top: 10px; left: 10px; z-index: 100; display: flex; flex-direction: column; gap: 5px; }}194            #debug-log {{ 195                position: absolute; bottom: 10px; left: 10px; 196                background: rgba(0,0,0,0.7); color: lime; 197                padding: 10px; font-size: 12px; pointer-events: none;198                max-height: 200px; overflow-y: auto; width: 400px;199                border-radius: 5px;200                display: block; 201            }}202            .btn {{ padding: 6px 10px; background: #333; color: white; border: 1px solid #555; cursor: pointer; border-radius: 3px; font-size: 12px; }}203            .btn:hover {{ background: #555; }}204            .row {{ display: flex; gap: 5px; align-items: center; }}205        </style>206        <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>207        <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script>208        <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/loaders/STLLoader.js"></script>209        <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/loaders/OBJLoader.js"></script>210        </head><body>211        212        <div id="ui">213            <div class="row">214                <button class="btn" onclick="fitCamera()">🔍 Auto Fit</button>215                <button class="btn" onclick="toggleLog()">📜 Log</button>216            </div>217            <div class="row" style="background:rgba(255,255,255,0.8); padding:5px; border-radius:3px; gap:10px;">218                <label><input type="checkbox" id="chkVisual" checked onchange="toggleVisual(this.checked)"> Show Visual</label>219                <label><input type="checkbox" id="chkCollision" onchange="toggleCollision(this.checked)"> Show Collision</label>220            </div>221        </div>222        <div id="debug-log">Initializing Renderer...</div>223 224        <script>225            function log(msg) {{226                const el = document.getElementById('debug-log');227                if(el) {{228                    el.innerHTML += "<div>> " + msg + "</div>";229                    el.scrollTop = el.scrollHeight;230                }}231                console.log(msg);232            }}233 234            function toggleLog() {{235                const el = document.getElementById('debug-log');236                el.style.display = (el.style.display === 'none') ? 'block' : 'none';237            }}238 239            const robotData = {json_data}; 240            const jointMeshes = {{}}; 241            const visualMeshes = [];      // [추가] 시각적 요소 리스트242            const collisionMeshes = [];243            const jointAxes = {{}};244            const rootGroup = new THREE.Group();245 246            const scene = new THREE.Scene(); scene.background = new THREE.Color(0xf0f0f0);247            248            // Camera (Z-Up)249            const camera = new THREE.PerspectiveCamera(50, window.innerWidth/window.innerHeight, 0.01, 10000);250            camera.up.set(0, 0, 1); 251            camera.position.set(2, 2, 2); 252            253            const renderer = new THREE.WebGLRenderer({{antialias:true}}); 254            renderer.setSize(window.innerWidth, window.innerHeight);255            renderer.shadowMap.enabled = true;256            document.body.appendChild(renderer.domElement);257            258            const controls = new THREE.OrbitControls(camera, renderer.domElement);259            260            scene.add(new THREE.AmbientLight(0xffffff, 0.6));261            const dl = new THREE.DirectionalLight(0xffffff, 0.8); 262            dl.position.set(5,10,7); scene.add(dl);263            264            const grid = new THREE.GridHelper(20, 20);265            grid.rotateX(Math.PI / 2); 266            scene.add(grid); 267            268            scene.add(new THREE.AxesHelper(1));269            scene.add(rootGroup);270 271            log("Renderer Started (V7).");272 273            function createThickFrame(len = 0.15, thick = 0.005) {{274                const group = new THREE.Group();275                const headLen = len * 0.2; 276                const headWidth = thick * 3; 277                278                function makeArrow(color, rot) {{279                    const arrow = new THREE.Group();280                    const mat = new THREE.MeshBasicMaterial({{color: color}});281                    const shaftGeo = new THREE.CylinderGeometry(thick, thick, len - headLen, 12);282                    const shaft = new THREE.Mesh(shaftGeo, mat);283                    shaft.position.y = (len - headLen) / 2;284                    arrow.add(shaft);285                    const headGeo = new THREE.ConeGeometry(headWidth, headLen, 12);286                    const head = new THREE.Mesh(headGeo, mat);287                    head.position.y = len - headLen / 2;288                    arrow.add(head);289                    arrow.rotation.set(...rot);290                    return arrow;291                }}292                group.add(makeArrow(0xff0000, [0, 0, -Math.PI/2]));293                group.add(makeArrow(0x00ff00, [0, 0, 0]));294                group.add(makeArrow(0x0000ff, [Math.PI/2, 0, 0]));295                return group;296            }}297 298            function createGeometry(info, isCollision=false) {{299                if (!info || !info.type) return null;300                let mesh = null;301 302                const mat = isCollision 303                    ? new THREE.MeshBasicMaterial({{color: 0xff0000, wireframe: true, transparent: true, opacity: 0.5}})304                    : new THREE.MeshPhongMaterial({{color: new THREE.Color(info.color[0], info.color[1], info.color[2]), shininess: 30, side: THREE.DoubleSide}});305 306                try {{307                    if(info.type === 'mesh' && info.mesh_data) {{308                        const loader = new THREE.STLLoader();309                        const binaryString = window.atob(info.mesh_data);310                        const len = binaryString.length;311                        const bytes = new Uint8Array(len);312                        for (let i = 0; i < len; i++) {{ bytes[i] = binaryString.charCodeAt(i); }}313                        const geo = loader.parse(bytes.buffer);314                        if(info.scale) geo.scale(info.scale[0], info.scale[1], info.scale[2]);315                        316                        geo.computeBoundingBox();317                        const size = new THREE.Vector3();318                        geo.boundingBox.getSize(size);319                        log("STL: " + size.x.toFixed(2) + " x " + size.y.toFixed(2) + " x " + size.z.toFixed(2));320                        321                        mesh = new THREE.Mesh(geo, mat);322                    }}323                    else if(info.type === 'obj' && info.mesh_data) {{324                        const loader = new THREE.OBJLoader();325                        const objGroup = loader.parse(info.mesh_data);326                        if(info.scale) objGroup.scale.set(info.scale[0], info.scale[1], info.scale[2]);327                        objGroup.traverse((child) => {{ if (child.isMesh) child.material = mat; }});328                        mesh = objGroup;329                        log("OBJ Loaded.");330                    }}331                    else if(info.type === 'cylinder') {{ 332                        const geo = new THREE.CylinderGeometry(info.dim[0], info.dim[0], info.dim[1], 32); 333                        geo.rotateX(Math.PI/2); 334                        mesh = new THREE.Mesh(geo, mat);335                    }}336                    else if(info.type === 'box') {{ 337                        const geo = new THREE.BoxGeometry(info.dim[0], info.dim[1], info.dim[2]); 338                        mesh = new THREE.Mesh(geo, mat);339                    }}340                    else if(info.type === 'sphere') {{ 341                        const geo = new THREE.SphereGeometry(info.dim[0], 32, 32); 342                        mesh = new THREE.Mesh(geo, mat);343                    }}344                    else if(info.type === 'error_box') {{345                        const geo = new THREE.BoxGeometry(0.2, 0.2, 0.2);346                        const errMat = new THREE.MeshBasicMaterial({{color: 0xff00ff, wireframe: true}});347                        mesh = new THREE.Mesh(geo, errMat);348                        log("⚠️ Fallback Box");349                    }}350                }} catch(e) {{351                    log("❌ Error: " + e.message);352                }}353                354                if(mesh && info.origin) {{ 355                    const [x,y,z, r,p,yaw] = info.origin; 356                    mesh.position.set(x,y,z); 357                    mesh.rotation.set(r,p,yaw, 'ZYX');358                }}359                360                // [수정] Visual/Collision 분리 저장361                if (isCollision && mesh) {{362                    mesh.visible = false; // 충돌체는 기본적으로 숨김363                    collisionMeshes.push(mesh);364                    if(mesh.isGroup) mesh.traverse(c => {{ if(c.isMesh) collisionMeshes.push(c); c.visible = false; }});365                }} 366                else if (!isCollision && mesh) {{367                    visualMeshes.push(mesh); // 시각체 리스트에 추가368                }}369                370                return mesh;371            }}372 373            function buildRobot(name, parent) {{374                const info = robotData.links[name]; 375                if(info) {{376                    if (info.visual) {{377                        const vMesh = createGeometry(info.visual, false);378                        if (vMesh) parent.add(vMesh);379                    }}380                    if (info.collision) {{381                        const cMesh = createGeometry(info.collision, true);382                        if (cMesh) parent.add(cMesh);383                    }}384                }}385                if(robotData.tree[name]) {{386                    robotData.tree[name].forEach(jName => {{387                        const jInfo = robotData.joints[jName];388                        const fixedGroup = new THREE.Group();389                        const [jx,jy,jz] = jInfo.xyz; 390                        const [jr,jp,jyaw] = jInfo.rpy;391                        fixedGroup.position.set(jx,jy,jz); 392                        fixedGroup.rotation.set(jr,jp,jyaw, 'ZYX');393                        parent.add(fixedGroup);394 395                        // [두꺼운 축 생성]396                        const axes = createThickFrame(0.15, 0.005);397                        axes.visible = false;398                        fixedGroup.add(axes);399                        jointAxes[jName] = axes;400 401                        const movingGroup = new THREE.Group();402                        fixedGroup.add(movingGroup);403                        jointMeshes[jName] = {{ group: movingGroup, axis: new THREE.Vector3(...jInfo.axis) }};404                        buildRobot(jInfo.child, movingGroup);405                    }});406                }}407            }}408 409            if(robotData.base) {{ buildRobot(robotData.base, rootGroup); }}410 411            window.fitCamera = function() {{412                const box = new THREE.Box3().setFromObject(rootGroup);413                if (box.isEmpty()) return;414                const size = box.getSize(new THREE.Vector3());415                const center = box.getCenter(new THREE.Vector3());416                if(size.length() < 0.0001) return;417                const maxDim = Math.max(size.x, size.y, size.z);418                const fov = camera.fov * (Math.PI / 180);419                let cameraZ = Math.abs(maxDim / 2 * Math.tan(fov * 2));420                cameraZ *= 2.5; 421                if (cameraZ < 0.1) cameraZ = 0.5;422                if (cameraZ > 1000) cameraZ = 1000;423                camera.position.set(center.x + cameraZ, center.y + cameraZ, center.z + cameraZ);424                camera.lookAt(center);425                controls.target.copy(center);426                controls.update();427                log("Camera fitted.");428            }};429 430            // [추가] 토글 기능들431            window.toggleVisual = function(isChecked) {{432                visualMeshes.forEach(mesh => {{ mesh.visible = isChecked; }});433            }};434 435            window.toggleCollision = function(isChecked) {{436                collisionMeshes.forEach(mesh => {{ mesh.visible = isChecked; }});437            }};438 439            function animate() {{ requestAnimationFrame(animate); renderer.render(scene, camera); }} 440            animate();441 442            setTimeout(fitCamera, 500);443            setTimeout(fitCamera, 1500);444 445            window.addEventListener("message", (e) => {{446                const data = e.data;447                if(Array.isArray(data)) {{448                    robotData.joint_order.forEach((name, i) => {{449                        if(jointMeshes[name]) {{450                            const rad = data[i] * (Math.PI/180);451                            const q = new THREE.Quaternion().setFromAxisAngle(jointMeshes[name].axis, rad);452                            jointMeshes[name].group.setRotationFromQuaternion(q);453                        }}454                    }});455                }}456                else if(data && data.type === 'frame') {{457                    if(jointAxes[data.name]) {{458                        jointAxes[data.name].visible = data.val;459                    }}460                }}461            }});462            window.addEventListener('resize', () => {{ 463                camera.aspect = window.innerWidth/window.innerHeight; 464                camera.updateProjectionMatrix(); 465                renderer.setSize(window.innerWidth, window.innerHeight); 466            }});467        </script></body></html>"""468        return html.escape(html_code)