gop checker

gop_checker.py

"""
Requires FFprobe installed.
"""
import argparse
import json
import logging
import subprocess


logger = logging.getLogger('encode_tests')
hdlr = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s: %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)


def probe_frames(file_, index=0):
    """ Yields a dict for each frame containing a bunch of information and frame count """
    process = subprocess.Popen("ffprobe -v 16 -show_frames -select_streams {} -print_format json=compact=1 {}".format(
        index, file_).split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1)
    logger.debug("Starting ffprobe.")
    for line in iter(process.stdout.readline, b''):
        if '"media_type"' in line:
            if line[-2:] == ",\n":
                line = line[:-2]
            elif line[-2:] == "\n":
                line = line[:-2]
            yield json.loads(line)


def keyframes(file_, index=0):
    for frame_count, line in enumerate(probe_frames(file_, index=0)):
        logger.debug("Processing frame {}".format(frame_count))
        if line['key_frame']:
            logger.debug("Keyframe found!")
            yield frame_count


def compare_sequences(sequences):
    match = None
    for file_, list_ in sequences.iteritems():
        if match:
            if match[1] != list_:
                raise AttributeError('{} keyframes does not match {} keyframes.'.format(match[0], file_))
        else:
            match = [file_, list_]


def run(sequence=False, compare=False, *files):
    if compare and len(files) < 2:
        raise RuntimeError("Can't compare with just one file.")
    sequences = {}
    for file_ in files:
        logger.info("Processing {}".format(file_))
        sequences[file_] = [c for c in keyframes(file_)]
    if sequence:
        for file_, list_ in sequences.iteritems():
            print "{}: {}".format(file_, list_)
    if compare:
        compare_sequences(sequences)
        print "Keyframes match across file set!"


def get_parser():
    parser = argparse.ArgumentParser(description='GOP analysing tools.')
    parser.add_argument(
        'files', metavar='video files.', nargs='+', help='The video stream must be the first in each file.')
    parser.add_argument(
        '--compare', dest='compare', action='store_true', default=False, help='Compare input files for a match.')
    parser.add_argument('--sequence', dest='sequence', action='store_true', default=False, help='Print GOP sequence.')
    parser.add_argument('--loglevel', default='info', choices=['error', 'warning', 'info', 'debug'])
    return parser


def get_args():
    return get_parser().parse_args()


def main():
    arguments = get_args()
    logger.setLevel(getattr(logging, arguments.loglevel.upper()))
    run(arguments.sequence, arguments.compare, *arguments.files)

if __name__ == '__main__':
    main()