Checking keyframe alignment using ffprobe

check_keyframes.py

import re
import subprocess
import logging
import math
import json
import pygal
import sys
logger = logging.getLogger(__name__)

'''
Usage from command line
./check_keyframes inputfile.mp4 2 25
'''

FFPROBE = "/usr/local/bin/ffprobe"

def get_keyframes(infile, keyframe_int, fps):
    # Should be able to return when keyframes are expected which is every x seconds
    # this would mean we should have an iframe every x * fps
    # Specific value for ffprobe output
    # TODO move this to using -print_format json to get results as a json object
    lines_past_type = 20
    line_count = 0
    print_line = False
    frame = 0
    found_iframe = False
    output_str = ""
    next_keyframe = math.ceil(float(fps) *  int(keyframe_int))
    key_framedist = next_keyframe

    cmd = FFPROBE + " -show_frames %s" % infile

    logger.info(cmd)
# -print_format json
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    while True:
      line = process.stdout.readline()
      if line != '':
        #the real code does filtering here
        #print line.rstrip()
        if (line.rstrip()=="media_type=video"):
            print_line = True
        if (print_line and (line_count < lines_past_type)):

                #print line.rstrip()
            if re.match(r'pict_type=I.*',line):
                found_iframe = True
                #print line.rstrip()
            match = re.match(r'coded_picture_number=(.*)',line)
            if (match):
                frame =  match.group(1)
                if (found_iframe):
                    output_str += "%s," % frame
                    found_iframe = False
                    if (int(frame)>=next_keyframe):
                        if (int(frame)==next_keyframe):
                            print "Found I frame as expected at %s" % next_keyframe
                        else:
                            print "Expected I frame at %s found at %s" % (next_keyframe, frame)
                        next_keyframe+=key_framedist

            line_count+=1
            if (line_count >= lines_past_type):
                print_line = False
                line_count = 0
      else:
        break
    return output_str
    
if __name__ == "__main__":
    get_keyframes(sys.argv[1], sys.argv[2], sys.argv[3])