Skip to content

Commit

Permalink
fixed slicing and exporting
Browse files Browse the repository at this point in the history
  • Loading branch information
Saeki-M committed Apr 25, 2023
1 parent 2eb9e68 commit 5ec5613
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 13 deletions.
31 changes: 19 additions & 12 deletions bvh.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@


class BvhNode:

def __init__(self, value=[], parent=None):
self.value = value
self.children = []
Expand Down Expand Up @@ -43,13 +42,21 @@ def name(self):


class Bvh:

def __init__(self, data):
self.data = data
self.root = BvhNode()
self.frames = []
self.tokenize()

@classmethod
def from_file(cls, filename):
with open(filename) as f:
mocap = cls(f.read())
return mocap

def __len__(self):
return round(self.nframes * 1000 / self.frame_rate)

def tokenize(self):
first_round = []
accumulator = ''
Expand Down Expand Up @@ -79,10 +86,13 @@ def tokenize(self):

def __getitem__(self, x):
if type(x) is int:
frames = self.frames[[round(x/(1000*self.frame_time))]]
frames = self.frames[[round(x/(1000*self.frame_rate))]]
elif type(x) is slice:
start_frame = round(x.start/(1000*self.frame_time))
end_frame = round(x.stop/(1000*self.frame_time))
start_time = x.start if x.start is not None else 0
end_time = x.stop if x.stop is not None else -1

start_frame = round(start_time/(1000*self.frame_rate))
end_frame = round(end_time/(1000*self.frame_rate))
frames = self.frames[start_frame:end_frame:x.step]
else:
raise KeyError
Expand Down Expand Up @@ -219,13 +229,10 @@ def joint_parent_index(self, name):

@property
def nframes(self):
try:
return int(next(self.root.filter('Frames:')).value[1])
except StopIteration:
raise LookupError('number of frames not found')
return len(self.frames)

@property
def frame_time(self):
def frame_rate(self):
try:
return float(next(self.root.filter('Frame')).value[2])
except StopIteration:
Expand All @@ -240,7 +247,7 @@ def raw_data(self):

data += "MOTION\n"
data += f"Frames:\t{self.nframes}\n"
data += f"Frame Time:\t{self.frame_time}\n"
data += f"Frame Time:\t{self.frame_rate}\n"

for frame in self.frames:
data += "\t".join(frame)+"\n"
Expand All @@ -262,7 +269,7 @@ def write_node(self, node, data, depth):
depth -= 1
return data, depth

def save(self, out_f):
def export(self, out_f):
with open(out_f, 'w') as f:
f.write(self.raw_data)

2 changes: 1 addition & 1 deletion tests/test_bvh.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ def test_nframes(self):
def test_frame_time(self):
with open('tests/test_freebvh.bvh') as f:
mocap = Bvh(f.read())
self.assertEqual(mocap.frame_time, 0.0333333)
self.assertEqual(mocap.frame_rate, 0.0333333)

def test_nframes2(self):
with open('tests/test_mocapbank.bvh') as f:
Expand Down

0 comments on commit 5ec5613

Please sign in to comment.