52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
import tables as tb
|
|
|
|
class Frame:
|
|
def __init__(self, id: str ='', frame_nr: int= 0, garment: str='', frame = None, frame_str: str = ''):
|
|
if frame is not None:
|
|
self.id: str = frame.id
|
|
self.frame_nr: int = frame.frame_nr
|
|
self.garment: str = frame.garment
|
|
assert id == '' and frame_nr == 0 and garment == '' and frame_str == ''
|
|
elif frame_str != '':
|
|
if ' ' in frame_str:
|
|
self.id: str = frame_str.split()[0]
|
|
self.frame_nr: int = int(frame_str.split()[1])
|
|
self.garment: str = frame_str.split()[2]
|
|
else:
|
|
self.id: str = frame_str.split('_')[0]
|
|
self.frame_nr: int = int(frame_str.split('_')[1])
|
|
self.garment: str = frame_str.split('_')[2]
|
|
assert id == '' and frame_nr == 0 and garment == '' and frame is None
|
|
else:
|
|
self.id: str = id
|
|
self.frame_nr: int = frame_nr
|
|
self.garment: str = garment
|
|
|
|
def __hash__(self):
|
|
return hash((self.id, self.frame_nr, self.garment))
|
|
|
|
def __eq__(self, other):
|
|
return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment)
|
|
|
|
def __str__(self):
|
|
return f'{self.id}_{self.frame_nr}_{self.garment}'
|
|
|
|
def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):
|
|
filters = tb.Filters(complevel=9, complib='blosc2')
|
|
group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
|
|
|
|
if type == 'float32':
|
|
atom = tb.Float32Atom()
|
|
elif type == 'uint8':
|
|
atom = tb.UInt8Atom()
|
|
elif type == 'bool':
|
|
atom = tb.BoolAtom()
|
|
else:
|
|
raise Exception('Unknown type')
|
|
|
|
if f'{name}' in group:
|
|
array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/{name}')
|
|
array[:] = data
|
|
else:
|
|
carray = h5file.create_carray(group, name, atom=atom, shape=data.shape, filters=filters)
|
|
carray[:] = data |