【python小工具】vbb文件提取json

参考 github caltech-pedestrian-converter

from scipy.io import loadmat
import os


def vbb2json(vbb_root):
    """逐级遍历，将所有vbb转换成json"""
    def _vbb2json(vbb_path):
        """转换单个vbb文件"""
        vbb = loadmat(vbb_path)['A'][0][0]

        nFrame = int(vbb[0][0][0])
        objLists = vbb[1][0]
        maxObj = int(vbb[2][0][0])
        objInit = vbb[3][0]
        objLbl = [str(v[0]) for v in vbb[4][0]]
        objStr = vbb[5][0]
        objEnd = vbb[6][0]
        objHide = vbb[7][0]
        altered = int(vbb[8][0][0])
        log = vbb[9][0]
        logLen = int(vbb[10][0][0])

        data = {}
        data['nFrame'] = nFrame
        data['maxObj'] = maxObj
        data['log'] = log.tolist()
        data['logLen'] = logLen
        data['altered'] = altered
        data['frames'] = defaultdict(list)

        for frame_id, obj in enumerate(objLists):
            if len(obj) > 0:
                for id, pos, occl, lock, posv in zip(obj['id'][0],
                                                     obj['pos'][0],
                                                     obj['occl'][0],
                                                     obj['lock'][0],
                                                     obj['posv'][0]):
                    keys = obj.dtype.names
                    id = int(id[0][0]) - 1  # MATLAB is 1-origin
                    p = pos[0].tolist()
                    pos = [p[0] - 1, p[1] - 1, p[2], p[3]]  # MATLAB is 1-origin
                    occl = int(occl[0][0])
                    lock = int(lock[0][0])
                    posv = posv[0].tolist()

                    datum = dict(zip(keys, [id, pos, occl, lock, posv]))
                    datum['lbl'] = str(objLbl[datum['id']])
                    # MATLAB is 1-origin
                    datum['str'] = int(objStr[datum['id']]) - 1
                    # MATLAB is 1-origin
                    datum['end'] = int(objEnd[datum['id']]) - 1
                    datum['hide'] = int(objHide[datum['id']])
                    datum['init'] = int(objInit[datum['id']])

                    data['frames'][frame_id].append(datum)

        with open(vbb_path[:-4]+'.json', 'w') as f:
            json.dump(data, f, sort_keys=True, indent=4)
        # print('实有 / 应有   %d / %d' % (len(data['frames']), data['nFrame']))

    # 遍历vbb_root下所有目录和文件
    for parent, _, filenames in os.walk(vbb_root):
        for filename in filenames:
            # vbb格式
            if filename[-4:] == '.vbb':
                vbb_path = os.path.join(parent, filename)
                _vbb2json(vbb_path)