Worm3DViewer/virtualworm.py at main · openworm/Worm3DViewer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import pyvista as pv

import sys

BODY_WALL_MUSCLE_OBJ_FILE = "bwm.obj"
NEURONS_OBJ_FILE = "neurons.obj"


def add_virtualworm_muscles(plotter, scale=20, translate=(-40, 0, 0)):
    print("Adding virtual worm file %s..." % BODY_WALL_MUSCLE_OBJ_FILE)
    mesh = pv.read(BODY_WALL_MUSCLE_OBJ_FILE)

    mesh.scale(scale, inplace=True)
    mesh.translate((translate), inplace=True)

    conn = mesh.connectivity("all")

    plotter.add_mesh(conn, smooth_shading=True, color="green")


def add_virtualworm_neurons(plotter, scale=20, translate=(-80, 0, 0)):
    print("Adding virtual worm file %s..." % NEURONS_OBJ_FILE)
    mesh2 = pv.read(NEURONS_OBJ_FILE)
    mesh2.scale(scale, inplace=True)
    mesh2.translate(translate, inplace=True)

    # plotter.add_mesh(mesh2, smooth_shading=True, color="orange")
    conn = mesh2.connectivity("all")
    plotter.add_mesh(
        conn,
        smooth_shading=True,
        cmap="jet",
    )
    plotter.remove_scalar_bar("RegionId")


if __name__ == "__main__":
    pl = pv.Plotter()

    add_virtualworm_muscles(pl)
    add_virtualworm_neurons(pl)

    pl.set_background("white")
    pl.add_axes()

    if "-nogui" not in sys.argv:
        pl.show()