scripts/csv2sto.py: Allow to specify muscle names on command line

1 files changed, 23 insertions, 8 deletions

diff --git a/scripts/csv2sto.py b/scripts/csv2sto.py
index 0f0e697..3a1b363 100755
--- a/scripts/csv2sto.py
+++ b/scripts/csv2sto.py
@@ -16,16 +16,18 @@ DEFAULT_NAME = 'control'
 DEFAULT_MAXTIME = 100.0
 
 def usage():
-    print """usage: %s [OPTION...] CSV-FILE STO-FILE
+    print("""usage: {} [OPTION...] CSV-FILE... STO-FILE
 
-  -f    force overwrite of existing files
-  -n    name of the sto file (default: %s)
-  -T N  scale time such that it goes from 0.0 to N
-  -h    show this help and exit""" % (DEFAULT_NAME, os.path.basename(sys.argv[0]))
+  -f        force overwrite of existing files
+  -m NAMES  use names from comma-separated list NAMES to name muscles instead
+            of CSV file name (# of names in list must match # of CSV files)
+  -n        name of the activation data in sto file (default: {})
+  -T N      scale time such that it goes from 0.0 to N
+  -h        show this help and exit""".format(DEFAULT_NAME, os.path.basename(sys.argv[0])))
 
 def main():
     try:
-        opts, args = getopt.getopt(sys.argv[1:], "fn:hT:")
+        opts, args = getopt.getopt(sys.argv[1:], "fm:n:hT:")
     except getopt.GetoptError, err:
         print(str(err))
         usage()
@@ -36,11 +38,14 @@ def main():
         sys.exit(-1)
 
     overwrite = False
+    muscles = None
     name = DEFAULT_NAME
     maxtime = DEFAULT_MAXTIME
     for o, a in opts:
         if o == '-f':
             overwrite = True
+        elif o == '-m':
+            muscles = [ x.strip() for x in a.split(',') ]
         elif o == '-n':
             name = a
         elif o == '-T':
@@ -56,6 +61,11 @@ def main():
         if not os.path.exists(f_csv):
             print "Error: CSV file %s does not exist" % f_csv
             sys.exit(-1)
+
+    if muscles and len(muscles) != len(fs_csv):
+        print("Error: number of names ({}) provided does not match number of CSV files ({})".format(len(muscles), len(fs_csv)))
+        sys.exit(-1)
+
     f_sto = args[-1]
     if not overwrite and os.path.exists(f_sto):
         print "Error: STO file %s already exists" % f_sto
@@ -95,9 +105,14 @@ def main():
     fd_sto.write("endheader\n")
 
     fd_sto.write("time")
+    i = 0
     for f_csv in fs_csv:
-        colname = os.path.splitext(os.path.basename(f_csv))[0]
-        print "[+] Adding muscle %s..." % colname
+        if muscles:
+            colname = muscles[i]
+            i += 1
+        else:
+            colname = os.path.splitext(os.path.basename(f_csv))[0]
+        print("[+] Adding muscle {}...".format(colname))
         fd_sto.write("\t" + colname)
     fd_sto.write("\n")