pad video sleap

pad_videos_slp.py

#!/usr/bin/env python3

import h5py
import cv2
import numpy as np
import sys
import json
import shutil

in_fname = sys.argv[1]
out_fname = sys.argv[2]

# in_fname = '/home/lili/Downloads/test-sleap/fly32.pkg.slp'
# out_fname = '/home/lili/Downloads/test-sleap/fly32_pad.pkg.slp'

# f = h5py.File(in_fname, 'r')

shutil.copy(in_fname, out_fname)

f = h5py.File(out_fname, 'r+')

for line in f['videos_json']:
    v = json.loads(line)
    name = v['backend']['dataset']
    print(name)

    if f[name].dtype in ['int8', 'uint8']:
        continue

    if len(f[name]) == 0:
        continue

    ds = f.pop(name)
    m = max([len(x) for x in ds])

    out = f.create_dataset(name, (len(ds), m),
                           compression='gzip',
                           dtype='int8', fillvalue=0)

    for i in range(len(ds)):
        x = ds[i]
        out[i, 0:len(x)] = x

    for k, v in ds.attrs.items():
        out.attrs[k] = v

f.close()

import tempfile
import subprocess
temp_filename = tempfile.NamedTemporaryFile().name
# temp_filename = temp_file.name

subprocess.run(['h5repack', out_fname, temp_filename])
shutil.move(temp_filename, out_fname)