summaryrefslogtreecommitdiff
path: root/scripts/osim-extract.py
blob: 1c0de6808e8d176651df4e185274ef8ddc573149 (plain)
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# osim-extract.py -- Extract information from .osim file
#
# Copyright (C) 2013 Tobias Klauser <tklauser@distanz.ch>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.

import os
import sys
import getopt
import xml.etree.ElementTree as ET

def usage():
    print("""usage: {} [OPTION...] FILE

options:
    -b  print only body names
    -j  print only joint names
    -m  print only muscle names
    -v  verbose mode
    -h  show this help and exit""".format(os.path.basename(sys.argv[0])))

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "bjmvh")
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit(-1)

    do_print = { 'b': False, 'j': False, 'm': False }
    verbose = False
    for o, a in opts:
        if o == '-b':
            do_print['b'] = True;
        elif o == '-j':
            do_print['j'] = True;
        elif o == '-m':
            do_print['m'] = True;
        elif o == '-v':
            verbose = True
        elif o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    if len(args) < 1:
        print("Error: no .osim model file specified")
        usage()
        sys.exit(-1)

    if not True in do_print.values():
        for k in do_print.iterkeys():
            do_print[k] = True

    try:
        tree = ET.parse(args[0])
    except:
        print("Error: failed to parse .osim model file")
        sys.exit(-1)

    root = tree.getroot()
    if verbose:
        print("Found root with version {}".format(root.attrib['Version']))
    model = tree.find('Model')
    if verbose:
        print("Found model '{}'".format(model.attrib['name']))

    if do_print['b']:
        bodies = [ body.attrib['name'] for body in root.findall('Model/BodySet/objects/Body') ]
        if verbose:
            print("bodies:")
        print(','.join(bodies))
    if do_print['j']:
        joints = [ joint.attrib['name'] for joint in root.findall('Model/BodySet/objects/Body/Joint/*') ]
        if verbose:
            print("joints:")
        print(','.join(joints))
    if do_print['m']:
        muscles = [ force.attrib['name'] for force in root.findall('Model/ForceSet/objects/Thelen2003Muscle') ]
        if verbose:
            print("muscles:")
        print(','.join(muscles))

if __name__ == '__main__':
    main()