Skip to content
Snippets Groups Projects
Select Git revision
  • 542a44ec9c59b34a5ed052c8c42b3a44bcbaef27
  • main default protected
2 results

ledchain_use.ino

Blame
  • TF2.tsx 9.21 KiB
    import { Props, useFrame, useThree } from '@react-three/fiber'
    import { PropsWithChildren, useEffect, useRef } from 'react';
    import { AxesHelper, Group, Matrix4, Mesh, MeshBasicMaterial, Object3D, Quaternion, SphereGeometry, Vector3 } from 'three'
    import { useConnection, useStore } from '../../common/store';
    import { Config as GeneralConfig } from '../../schemas/Config.schema';
    import tf2_msgs from '../../common/ROS/tf2_msgs';
    import geometry_msgs from '../../common/ROS/geometry_msgs';
    import { Topic } from 'roslib';
    import { load as load_yaml } from 'js-yaml'
    import { Text } from 'troika-three-text'
    
    const prefix = "tf2_tree_frame_";
    
    export interface FrameProps {
        frameId?: string;
    }
    
    export function Frame(props: PropsWithChildren<FrameProps>) {
        const { scene } = useThree();
        const group = useRef<Group | null>(null);
    
        (window as any).scene = scene;
    
        useEffect(() => {
            if (!props.frameId || props.frameId == "") return;
            const timer = setInterval(() => {
                const frameObject = scene.getObjectByProperty("name", prefix + props.frameId);
                if (frameObject && group.current) {
                    frameObject.add(group.current);
                    clearInterval(timer);
                }
            }, 1000);
            return () => { clearInterval(timer); };
        }, [group, scene, props.frameId]);
    
        return (
            <group ref={group}>
                {props.children}
            </group>
        );
    }
    
    export type yaml_anchor = {
        frame_id: string;
        p: [number, number, number]; //[x,y,z] in ROS
        q: [number, number, number, number]; //[w,x,y,z] in ROS
    };
    
    export interface TF2Props {
        yaml?: string[];
        attachPrintGraphFunction?: boolean;
    }
    
    function convertFromYAML(translation: [number, number, number], rotation: [number, number, number, number]): Matrix4 {
        const position: Vector3 = new Vector3(-translation[0], translation[2], translation[1]);
        const orientation: Quaternion = new Quaternion(-rotation[1], rotation[3], rotation[2], rotation[0]);
        return new Matrix4().compose(position, orientation, new Vector3(1, 1, 1));
    }
    
    function convertFromROS(poseStamped: geometry_msgs["Transform"]): Matrix4 {
        const position: Vector3 = new Vector3(-poseStamped.translation.x, poseStamped.translation.z, poseStamped.translation.y);
        const orientation: Quaternion = new Quaternion(-poseStamped.rotation.x, poseStamped.rotation.z, poseStamped.rotation.y, poseStamped.rotation.w);
        return new Matrix4().compose(position, orientation, new Vector3(1, 1, 1));
    }
    
    function TF2(props: TF2Props) {
        const { scene } = useThree();
        let root = useRef<Object3D>(new Object3D());
        let knownObjects = useRef(new Map<string, Object3D>());
        let debugGeometry = useRef<Object3D[]>([]);
        let debugLabels = useRef<Object3D[]>([]);
        const debug_geometry = useRef<SphereGeometry>(new SphereGeometry(0.1, 16, 16));
        const debug_material = useRef<MeshBasicMaterial>(new MeshBasicMaterial({ color: 0xff0000 }));
        const config = useStore(state => state.transformTree);
        const connection = useConnection(config?.rosbridge.uri);
        let showDebugSpheres = useStore(state => state.appSettings.showTFDebugVis);
        let showDebugLabels = useStore(state => state.appSettings.showTFDebugLabels);
    
        useEffect(() => {
            if (!props.yaml) return;
            for (let i = 0; i < props.yaml.length; i++) {
                fetch(props.yaml[i])
                    .then((r) => {
                        return r.text();
                    }).then((yaml) => {
                        let doc = load_yaml(yaml) as { [key: string]: yaml_anchor };
                        Object.entries(doc).map(([to_id, node]) => {
                            if (!knownObjects.current.has(node.frame_id)) createFrame(node.frame_id);
                            if (!knownObjects.current.has(to_id)) createFrame(to_id, node.frame_id);
    
                            const to_frame = knownObjects.current.get(to_id);
    
                            if (to_frame) to_frame.matrix = convertFromYAML(node.p, node.q);
    
                            console.log("YAML: attaching " + to_id + " to " + node.frame_id);
                        });
                    });
            }
        }, [props.yaml]);
    
        useEffect(
            () => {
                root.current.name = "tf2_root";
                scene.add(root.current);
                return () => { scene.remove(root.current); };
            }, [scene]
        );
    
        function createFrame(name: string, parent?: string): Object3D {
            const Result = new Object3D();
            Result.name = prefix + name;
            Result.matrixAutoUpdate = false;
    
            // Debug geometry
            const axesHelper = new AxesHelper();
            axesHelper.visible = showDebugSpheres;
            debugGeometry.current.push(axesHelper);
            Result.add(axesHelper);
    
            // const sphere = new Mesh(debug_geometry.current, debug_material.current);
            // sphere.visible = showDebugSpheres;
            // debugGeometry.current.push(sphere);
            // Result.add(sphere);
    
            // Debug Label
            const label = new Text();
            label.visible = true;
            label.text = name;
            label.fontSize = 0.2;
            label.anchorX = "center";
            label.anchorY = "bottom";
            label.color = 0xFF0000;
            debugLabels.current.push(label);
    
            Result.add(label);
    
            knownObjects.current.set(name, Result);
            if (parent) {
                knownObjects.current.get(parent)?.add(Result); // Attach to parent
            } else {
                root.current.add(Result);
            }
            return Result;
        }
    
        function newTFMessage(message: ROSLIB.Message) {
            const parsedMessage = message as tf2_msgs["TFMessage"];
    
            for (let i = 0; i < parsedMessage.transforms.length; i++) {
                const transform = convertFromROS(parsedMessage.transforms[i].transform);
                const from_id = parsedMessage.transforms[i].header.frame_id; //parent
                const to_id = parsedMessage.transforms[i].child_frame_id; //child
    
                //Create frames if not known
                if (!knownObjects.current.has(from_id)) createFrame(from_id);
                if (!knownObjects.current.has(to_id)) createFrame(to_id, from_id);
    
                const from_frame = knownObjects.current.get(from_id);
                const to_frame = knownObjects.current.get(to_id);
    
                if (!from_frame || !to_frame) continue; //something went wrong
    
                //Attach child to parent if not already done
                if (to_frame.parent !== from_frame) {
                    from_frame.add(to_frame);
                }
    
                //Ignore UTM frame, as it is possibly far far away
                if (from_id == "utm" || from_id.endsWith("/utm")) {
                    to_frame.matrix = new Matrix4();
                } else {
                    to_frame.matrix = transform;
                }
            }
        }
    
        // printGraph() function
        useEffect(() => {
            if (!(window as any).printGraph && (props.attachPrintGraphFunction || false)) {
                (window as any).printGraph = () => {
                    let recurse = function (node: Object3D): string {
                        let node_name = node.name;
                        if (node_name.startsWith(prefix)) node_name = node_name.slice(prefix.length);
                        let res = "\"" + node_name + "\" : {";
                        let first = true;
                        for (let i = 0; i < node.children.length; i++) {
                            if (node.children[i].name === "" || debugGeometry.current.indexOf(node) != -1) continue;
                            if (first) { first = false; } else { res += ","; }
                            res += recurse(node.children[i]);
                        }
                        res += "}";
                        return res;
                    };
                    return console.log(JSON.stringify(JSON.parse("{" + recurse(root.current) + "}"), null, 2));
                };
            }
            return () => { (window as any).printGraph = undefined };
        }, [root.current, props.attachPrintGraphFunction]);
    
        // visualizations visible/invisible
        useEffect(() => {
            for (let obj of debugGeometry.current) {
                obj.visible = showDebugSpheres;
            }
    
            for (let label of debugLabels.current) {
                label.visible = showDebugLabels;
            }
        }, [showDebugSpheres, showDebugLabels]);
    
        useFrame(({ camera }) => {
            if (!showDebugLabels) return;
    
            for (let label of debugLabels.current) {
                if (!label.parent) continue;
    
                // always face the camera
                let q = new Quaternion();
                let p = new Quaternion();
                camera.getWorldQuaternion(q);
                label.parent.getWorldQuaternion(p);
                label.setRotationFromQuaternion(p.invert().multiply(q));
            }
        })
    
        // Subscribe/Unsubscribe to TF Messages
        useEffect(() => {
            if (!connection?.rosbridge || !config) return;
    
            let tf_topic = new Topic({ "ros": connection?.rosbridge, "name": "/tf", "messageType": 'tf2_msgs/TFMessage' });
            let tf_static_topic = new Topic({ "ros": connection?.rosbridge, "name": "/tf_static", "messageType": 'tf2_msgs/TFMessage' });
    
            tf_topic.subscribe(newTFMessage);
            tf_static_topic.subscribe(newTFMessage);
    
            return () => {
                tf_topic.unsubscribe(newTFMessage);
                tf_static_topic.unsubscribe(newTFMessage);
            };
        }, [connection?.rosbridge, config]);
    
        return null;
    }
    
    export default TF2;