make_mesh_png.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
import argparse
from glob import glob
from pathlib import Path
 
import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv
 
 
def expand_input_files(inputs):
    files = []
    for entry in inputs:
        matches = sorted(glob(entry))
        if matches:
            files.extend(matches)
        else:
            files.append(entry)
    return files
 
 
def filter_vtk_files(files):
    return [file for file in files if file.lower().endswith(".vtk")]
 
 
def select_field_component(mesh, field_name, component):
    if field_name in mesh.point_data:
        data = mesh.point_data[field_name]
        target = mesh.point_data
    elif field_name in mesh.cell_data:
        data = mesh.cell_data[field_name]
        target = mesh.cell_data
    else:
        raise KeyError(f"Field '{field_name}' was not found in point or cell data.")
 
    values = np.asarray(data)
    if values.ndim == 1:
        values = values.reshape(-1, 1)
    else:
        values = values.reshape(values.shape[0], -1)
 
    component_index = component
 
    if component_index < 0 or component_index >= values.shape[1]:
        raise ValueError(
            f"Field '{field_name}' has {values.shape[1]} component(s), "
            f"so component {component} is out of range. "
            f"Use 0..{values.shape[1] - 1}."
        )
 
    component_name = f"{field_name}[{component_index}]"
    target[component_name] = values[:, component_index]
    return component_name
 
 
def make_png(file, args):
    my_cmap = plt.cm.get_cmap("turbo", 124)
    mesh = pv.read(file)
    output_file = str(Path(file).with_suffix(".png"))
 
    if args.wrap_vector:
        mesh = mesh.warp_by_vector(args.wrap_vector, factor=args.wrap_factor)
 
    plotter = pv.Plotter(notebook=False, off_screen=True)
    try:
        if args.field:
            scalar_name = args.field
            if args.field_component is not None:
                scalar_name = select_field_component(
                    mesh, args.field, args.field_component
                )
            plotter.add_mesh(
                mesh,
                scalars=scalar_name,
                show_edges=args.show_edges,
                smooth_shading=False,
                cmap=my_cmap,
            )
        else:
            plotter.add_mesh(
                mesh,
                show_edges=args.show_edges,
                edge_color="white",
                color="white",
            )
        if args.d2:
            plotter.camera_position = args.d2
        plotter.camera.zoom(args.zoom)
        plotter.camera.roll += args.roll
        plotter.camera.azimuth += args.azimuth
        plotter.screenshot(output_file)
    finally:
        plotter.close()
 
 
def maybe_start_virtual_display():
    try:
        from pyvirtualdisplay import Display
    except ImportError:
        return None
 
    try:
        display = Display(backend="xvfb", visible=False, size=(800, 800))
        display.start()
    except Exception:
        return None
    return display
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert multiple vtk files to png files using PyVista."
    )
    parser.add_argument(
        "files",
        help="list of vtk files or glob masks",
        nargs="+",
    )
    parser.add_argument("-d2", "--d2", dest="d2", default="")
    parser.add_argument("-f", "--field", dest="field", default="", type=str)
    parser.add_argument(
        "--field-component",
        dest="field_component",
        default=None,
        type=int,
        help="0-based component index for vector/tensor fields",
    )
    parser.add_argument("-wv", "--wrap_vector", dest="wrap_vector", default="", type=str)
    parser.add_argument("--wrap-factor", dest="wrap_factor", default=1.0, type=float)
    parser.add_argument("--zoom", dest="zoom", default=1.2, type=float)
    parser.add_argument("--roll", dest="roll", default=0, type=float)
    parser.add_argument("--azimuth", dest="azimuth", default=0, type=float)
    parser.add_argument("--show-edges", dest="show_edges", action="store_true", default=True, help="Show mesh edges (default)")
    parser.add_argument("--no-show-edges", dest="show_edges", action="store_false", help="Hide mesh edges")
    parser.add_argument("--debug", dest="debug", action="store_true")
    args = parser.parse_args()
 
    if args.debug:
        print(args)
 
    file_list = filter_vtk_files(expand_input_files(args.files))
    if not file_list:
        parser.error("No vtk files were found with the given names or glob masks.")
 
    display = maybe_start_virtual_display()
    try:
        for file_name in file_list:
            make_png(file_name, args)
    finally:
        if display is not None:
            display.stop()