io_anim_c3d/import_c3d.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 3
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8-80 compliant>
  20
  21 # By Daniel Monteiro Basso, April-November 2011.
  22
  23 # This script was developed with financial support from the Foundation for
  24 # Science and Technology of Portugal, under the grant SFRH/BD/66452/2009.
  25
  26 # Complete rewrite, but based on the original importer for Blender
  27 # 2.39, developed by Jean-Baptiste PERIN (jb_perin(at)yahoo.fr), which was
  28 # based on the MATLAB C3D loader from Alan Morris, Toronto, October 1998
  29 # and Jaap Harlaar, Amsterdam, april 2002
  30
  31
  32 import struct
  33 try:
  34     from numpy import array as vec  # would be nice to have NumPy in Blender
  35 except:
  36     from mathutils import Vector as vec
  37
  38
  39 class Marker:
  40     position = (0., 0., 0.)
  41     confidence = -1.
  42
  43
  44 class Parameter:
  45     def __init__(self, infile):
  46         (nameLength, self.paramIdx) = struct.unpack('bb', infile.read(2))
  47         if not nameLength:
  48             self.name = ''
  49             return
  50         nameLength = abs(nameLength)  # negative flags something
  51         if nameLength > 64:
  52             raise ValueError
  53         self.name = infile.read(nameLength).decode('ascii')
  54         (offset, b) = struct.unpack('hb', infile.read(3))
  55         if self.paramIdx > 0:
  56             self.isGroup = False
  57             self.data = infile.read(offset - 3)
  58         else:
  59             self.isGroup = True
  60             self.paramIdx *= -1
  61             self.description = infile.read(b)
  62             self.params = {}
  63
  64     def collect(self, infile):
  65         while True:
  66             p = Parameter(infile)
  67             if not p.name or p.isGroup:
  68                 return p
  69             self.params[p.name] = p
  70
  71     def decode(self):
  72         # for now only decode labels
  73         l, c = struct.unpack('BB', self.data[1:3])
  74         return [self.data[3 + i:3 + i + l].strip().decode('ascii')
  75                 for i in range(0, l * c, l)]
  76
  77
  78 class MarkerSet:
  79     def __init__(self, fileName, scale=1., stripPrefix=True, onlyHeader=False):
  80         self.fileName = fileName
  81         if fileName.endswith('.csv'):
  82             with open(fileName, 'rt') as infile:
  83                 self.readCSV(infile)
  84             return
  85         if onlyHeader:
  86             self.infile = open(fileName, 'rb')
  87             self.readHeader(self.infile, scale)
  88             self.identifyMarkerPrefix(stripPrefix)
  89             self.infile.seek(512 * (self.dataBlock - 1))
  90             self.frames = []
  91             return
  92         with open(fileName, 'rb') as infile:
  93             self.readHeader(infile, scale)
  94             self.identifyMarkerPrefix(stripPrefix)
  95             self.readFrameData(infile)
  96
  97     def readCSV(self, infile):
  98         import csv
  99         csvr = csv.reader(infile)
 100         header = next(csvr)
 101         if 0 != len(header) % 3:
 102             raise Exception('Incorrect data format in CSV file')
 103         self.markerLabels = [label[:-2] for label in header[::3]]
 104         self.frames = []
 105         for framerow in csvr:
 106             newFrame = []
 107             for c in range(0, len(framerow), 3):
 108                 m = Marker()
 109                 try:
 110                     m.position = vec([float(v) for v in framerow[c:c + 3]])
 111                     m.confidence = 1.
 112                 except:
 113                     pass
 114                 newFrame.append(m)
 115             self.frames.append(newFrame)
 116         self.startFrame = 0
 117         self.endFrame = len(self.frames) - 1
 118         self.scale = 1.
 119
 120     def writeCSV(self, fileName, applyScale=True, mfilter=[]):
 121         import csv
 122         with open(fileName, 'w') as fo:
 123             o = csv.writer(fo)
 124             appxyz = lambda m: [m + a for a in ('_X', '_Y', '_Z')]
 125             explabels = (appxyz(m) for m in self.markerLabels
 126                          if not mfilter or m in mfilter)
 127             o.writerow(sum(explabels, []))
 128             fmt = lambda m: tuple('{0:.4f}'.format(
 129                 a * (self.scale if applyScale else 1.))
 130                 for a in m.position)
 131             nan = ('NaN', 'NaN', 'NaN')
 132             if mfilter:
 133                 mfilter = [self.markerLabels.index(m)
 134                             for m in self.markerLabels if m in mfilter]
 135             for f in self.frames:
 136                 F = f
 137                 if mfilter:
 138                     F = [m for i, m in enumerate(f) if i in mfilter]
 139                 expmarkers = (m.confidence < 0 and nan or fmt(m) for m in F)
 140                 o.writerow(sum(expmarkers, ()))
 141
 142     def identifyMarkerPrefix(self, stripPrefix):
 143         prefix = self.markerLabels[0]
 144         for ml in self.markerLabels[1:]:
 145             if len(ml) < len(prefix):
 146                 prefix = prefix[:len(ml)]
 147             if not prefix:
 148                 break
 149             for i in range(len(prefix)):
 150                 if prefix[i] != ml[i]:
 151                     prefix = prefix[:i]
 152                     break
 153         self.prefix = prefix
 154         if stripPrefix:
 155             p = len(self.prefix)
 156             self.markerLabels = [ml[p:] for ml in self.markerLabels]
 157
 158     def readHeader(self, infile, scale):
 159         (self.firstParameterBlock, key, self.markerCount, bogus,
 160          self.startFrame, self.endFrame,
 161          bogus) = struct.unpack('BBhhhhh', infile.read(12))
 162         if key != 80:
 163             raise Exception('Not a C3D file.')
 164         self.readParameters(infile)
 165         infile.seek(12)
 166         td = infile.read(12)
 167         if self.procType == 2:
 168             td = td[2:4] + td[:2] + td[4:8] + td[10:] + td[8:10]
 169         (self.scale, self.dataBlock, bogus,
 170          self.frameRate) = struct.unpack('fhhf', td)
 171         self.scale *= scale
 172         if self.scale < 0:
 173             if self.procType == 2:
 174                 self.readMarker = self.readFloatMarkerInvOrd
 175             else:
 176                 self.readMarker = self.readFloatMarker
 177             self.scale *= -1
 178         else:
 179             self.readMarker = self.readShortMarker
 180
 181     def readParameters(self, infile):
 182         infile.seek(512 * (self.firstParameterBlock - 1))
 183         (ig, ig, pointIdx,
 184          self.procType) = struct.unpack('BBBB', infile.read(4))
 185         self.procType -= 83
 186         if self.procType not in {1, 2}:
 187             # 1(INTEL-PC); 2(DEC-VAX); 3(MIPS-SUN/SGI)
 188             print('Warning: importer was not tested for files from '
 189                   'architectures other than Intel-PC and DEC-VAX')
 190             print('Type: {0}'.format(self.procType))
 191         self.paramGroups = {}
 192         g = Parameter(infile)
 193         self.paramGroups[g.name] = g
 194         while(True):
 195             g = g.collect(infile)
 196             if not g.name:
 197                 break
 198             self.paramGroups[g.name] = g
 199         cand_mlabel = None
 200         for pg in self.paramGroups:
 201             #print("group: " + pg)
 202             #for p in self.paramGroups[pg].params:
 203             #    print("   * " + p)
 204             if 'LABELS' in self.paramGroups[pg].params:
 205                 cand_mlabel = self.paramGroups[pg].params['LABELS'].decode()
 206                 if len(cand_mlabel) == self.markerCount:
 207                     break
 208                 cand_mlabel = None
 209         # pg should be 'POINT', but let's be liberal and accept any group
 210         # as long as the LABELS parameter has the same number of markers
 211         if cand_mlabel is None:
 212             self.markerLabels = ["m{}".format(idx)
 213                                  for idx in range(self.markerCount)]
 214         else:
 215             self.markerLabels = cand_mlabel
 216         repeats = {}
 217         for i, m in enumerate(self.markerLabels):
 218             if m in repeats:
 219                 self.markerLabels[i] = '{}.{}'.format(m, repeats[m])
 220                 repeats[m] += 1
 221             else:
 222                 repeats[m] = 1
 223
 224     def readMarker(self, infile):
 225         pass  # ...
 226
 227     def readFloatMarker(self, infile):
 228         m = Marker()
 229         x, y, z, m.confidence = struct.unpack('ffff', infile.read(16))
 230         m.position = (x * self.scale, y * self.scale, z * self.scale)
 231         return m
 232
 233     def readFloatMarkerInvOrd(self, infile):
 234         m = Marker()
 235         inv = lambda f: f[2:] + f[:2]
 236         i = lambda: inv(infile.read(4))
 237         x, y, z, m.confidence = struct.unpack('ffff', i() + i() + i() + i())
 238         m.position = (x * self.scale, y * self.scale, z * self.scale)
 239         return m
 240
 241     def readShortMarker(self, infile):
 242         m = Marker()
 243         x, y, z, m.confidence = struct.unpack('hhhh', infile.read(8))
 244         m.position = (x * self.scale, y * self.scale, z * self.scale)
 245         return m
 246
 247     def readFrameData(self, infile):
 248         infile.seek(512 * (self.dataBlock - 1))
 249         self.frames = []
 250         for f in range(self.startFrame, self.endFrame + 1):
 251             frame = [self.readMarker(infile) for m in range(self.markerCount)]
 252             self.frames.append(frame)
 253
 254     def readNextFrameData(self):
 255         if len(self.frames) < (self.endFrame - self.startFrame + 1):
 256             frame = [self.readMarker(self.infile)
 257                 for m in range(self.markerCount)]
 258             self.frames.append(frame)
 259         return self.frames[-1]
 260
 261     def getFramesByMarker(self, marker):
 262         if type(marker) == int:
 263             idx = marker
 264         else:
 265             idx = self.markerLabels.index(marker)
 266         fcnt = self.endFrame - self.startFrame + 1
 267         return [self.frames[f][idx] for f in range(fcnt)]
 268
 269     def getMarker(self, marker, frame):
 270         idx = self.markerLabels.index(marker)
 271         return self.frames[frame - self.startFrame][idx]
 272
 273
 274 def read(filename, *a, **kw):
 275     return MarkerSet(filename, *a, **kw)
 276
 277 # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
 278
 279
 280 if __name__ == '__main__':
 281     import os
 282     import sys
 283
 284     sys.argv.pop(0)
 285     if not sys.argv:
 286         print("Convert C3D to CSV.\n"
 287               "Please specify at least one C3D input file.")
 288         raise SystemExit
 289     while sys.argv:
 290         fname = sys.argv.pop(0)
 291         markerset = read(fname)
 292         print("frameRate={0.frameRate}\t"
 293               "scale={0.scale:.2f}\t"
 294               "markers={0.markerCount}\t"
 295               "startFrame={0.startFrame}\t"
 296               "endFrame={0.endFrame}".format(markerset))
 297         markerset.writeCSV(fname.lower().replace(".c3d", ".csv"))