Takes mediapipe results and converts them to a collada file.

mediapipe_draw_collada_utils.py

import collada
import numpy as np


class DrawCollada:
    def __init__(self):
        self.mesh = collada.Collada()
        self.scene_nodes = []
        self.counter = 0

    def draw_landmarks(self, landmark_list, connections):
        """
        Takes mediapipe results and converts them to a collada file.
        Handy for debugging.
        """
        self.counter += 1
        c = str(self.counter)
        vert_floats = []
        for pos in landmark_list.landmark:
            vert_floats.append(pos.x)
            vert_floats.append(pos.y)
            vert_floats.append(pos.z)
        indices_lines = []
        for conn in connections:
            indices_lines.append(conn[0])
            indices_lines.append(conn[1])

        vert_src = collada.source.FloatSource(
            "verts-array" + c, np.array(vert_floats), ("X", "Y", "Z")
        )

        geom = collada.geometry.Geometry(self.mesh, "geom" + c, "geom" + c, [vert_src])

        input_list = collada.source.InputList()
        input_list.addInput(0, "VERTEX", "#" + vert_src.id)

        lines = geom.createLineSet(
            np.array(indices_lines), input_list, "fakeMaterialid"
        )
        geom.primitives.append(lines)

        self.mesh.geometries.append(geom)
        geomnode = collada.scene.GeometryNode(geom)
        node = collada.scene.Node("node" + c, children=[geomnode])
        self.scene_nodes.append(node)

    def write_scene(self, path):
        myscene = collada.scene.Scene("myscene", self.scene_nodes)
        self.mesh.scenes.append(myscene)
        self.mesh.scene = myscene

        self.mesh.write(path)