utils.py

import os
import glob
import json
from pathlib import Path
from collections import OrderedDict
import ffmpeg


def read_json(fname):
    fname = Path(fname)
    with fname.open('rt') as handle:
        return json.load(handle, object_hook=OrderedDict)


def get_configs(base_path, filename_pattern="*.json"):
    configs = []
    glob_pattern = os.path.join(base_path, filename_pattern)
    file_paths = sorted(glob.glob(glob_pattern))
    for file_path in file_paths:
        file_name = Path(file_path).stem
        config = read_json(file_path)
        config['name'] = file_name
        configs.append(config)
    return configs


def get_display_names(configs):
    display_names = []
    for config in configs:
        display_names.append(config['display_name'])
    return display_names


def get_path_for_viz(basedir, dataset, sequence, model, viz):
    dat_seq = "_".join([dataset["name"], sequence])
    video_path = os.path.join(basedir, model['model_id'], dat_seq, "videos", viz["name"] + ".mp4")
    return video_path


def get_text_str(height, width, text, fontsize, fontcolor="black"):
    return "drawtext=text='{}':" \
           "fontsize={}:fontcolor={}:" \
           "x=({}-text_w)/2:y=({}-text_h)/2".format(text, fontsize, fontcolor, width, height)


def get_video_height(video_path):
    probe = ffmpeg.probe(video_path)
    video_streams = ffmpeg.probe(video_path, select_streams="v")
    return video_streams['streams'][0]['height']