lattice_infill.py

"""Example script demonstrating filling a triangle mesh with a gyroid.

Usage:
    python examples/lattice_infill.py -t <token> -p <project> examples/suzanne.obj suzanne_gyroid.stl

For more details on the available job types please refer to the Metafold REST API
documentation.
"""
from argparse import ArgumentParser
from math import sqrt
from metafold import MetafoldClient
from typing import Any, TypeAlias, TypedDict
import argparse
import os
import sys

Vec3: TypeAlias = [float, float, float]


class Patch(TypedDict):
    offset: Vec3
    size: Vec3
    resolution: Vec3


def main() -> None:
    parser = ArgumentParser(description="Fill mesh with gyroid")
    parser.add_argument("-t", "--token", type=str, help="access token")
    parser.add_argument("-p", "--project", type=str, help="project id", required=True)
    parser.add_argument("infile", type=argparse.FileType("rb"), help="mesh file")
    parser.add_argument("outfile", type=str, help="output file", nargs="?")

    args = parser.parse_args()

    token = args.token or os.environ.get("METAFOLD_ACCESS_TOKEN")
    if not token:
        parser.error("access token is required")

    metafold = MetafoldClient(token, args.project)

    # Assets must be unique by filename. We query the project assets for an asset with a
    # matching filename and update the existing one accordingly, otherwise we simply
    # create a new asset.
    print("Uploading mesh asset...")
    filename = os.path.basename(args.infile.name)
    if existing := metafold.assets.list(q=f"filename:{filename}"):
        mesh = metafold.assets.update(existing[0].id, args.infile)
    else:
        mesh = metafold.assets.create(args.infile)

    # Run a "sample_triangle_mesh" job to generate a volume from the triangle mesh
    print("Running sample_triangle_mesh job...")
    sample_mesh_job = metafold.jobs.run("sample_triangle_mesh", {
        "mesh_filename": mesh.filename,
        "max_resolution": 256,
    })

    # "sample_triangle_mesh" generates multiple assets. We recommend using the filename
    # extension to find the generated asset of interest.
    for asset in sample_mesh_job.assets:
        if asset.filename.endswith(".bin"):
            volume_filename = asset.filename
            break
    else:
        print("Failed to generate mesh volume", file=sys.stderr)
        sys.exit(1)

    # Some jobs have JSON metadata which is automatically decoded into a dictionary
    patch = sample_mesh_job.meta["patch"]

    # We define a Metafold geometry graph that loads the volume generated by the
    # "sample_triangle_mesh" job and performs a CSG intersection with a gyroid.
    graph = create_graph(volume_filename, patch)

    # Run a "export_triangle_mesh" job to evaluate the geometry graph and generate a
    # tessellated mesh asset.
    print("Running export_triangle_mesh job...")
    export_job = metafold.jobs.run("export_triangle_mesh", {
        "graph": graph,
        "point_source": 0,
    })

    if args.outfile:
        print("Downloading generated mesh asset...")
        export_asset = export_job.assets[0].id
        metafold.assets.download_file(export_asset, args.outfile)


def create_graph(volume_filename: str, patch: Patch) -> dict[str, Any]:
    # SDF narrow-band width is relative to grid cell size
    cell_size = [
        size / (resolution - 1) for
        size, resolution in zip(patch["size"], patch["resolution"])
    ]
    threshold_width = 3.0 * sqrt(
          cell_size[0] ** 2
        + cell_size[1] ** 2
        + cell_size[2] ** 2
    )

    return {
        "operators": [
            {
                "type": "GenerateSamplePoints",
                "parameters": patch,
            },
            {
                "type": "LoadVolume",
                "parameters": {
                    "volume_data": {
                        "file_type": "Raw",
                        "path": volume_filename,
                    },
                    "resolution": patch["resolution"],
                },
            },
            {
                "type": "SampleVolume",
                "parameters": {
                    "volume_offset": patch["offset"],
                    "volume_size": patch["size"],
                },
            },
            {
                "type": "SampleSurfaceLattice",
                "parameters": {
                    "lattice_type": "Gyroid",
                    "scale": [10.0, 10.0, 10.0],
                },
            },
            {
                "type": "CSG",
                "parameters": {
                    "operation": "Intersect",
                },
            },
            {
                "type": "Redistance",
                "parameters": {
                    "size": patch["size"],
                }
            },
            {
                "type": "Threshold",
                "parameters": {
                    "width": threshold_width,
                }
            },
        ],
        "edges": [
            {"source": 0, "target": [2, "Points"]},   # GenerateSamplePoints -> SampleVolume
            {"source": 1, "target": [2, "Volume"]},   # LoadVolume -> SampleVolume
            {"source": 2, "target": [4, "A"]},        # SampleVolume -> CSG
            {"source": 0, "target": [3, "Points"]},   # GenerateSamplePoints -> SampleSurfaceLattice
            {"source": 3, "target": [4, "B"]},        # SampleSurfaceLattice -> CSG
            {"source": 4, "target": [5, "Samples"]},  # CSG -> Redistance
            {"source": 5, "target": [6, "Samples"]},  # Redistance -> Threshold
        ],
    }


if __name__ == "__main__":
    main()