segmenter

#!/usr/bin/python
import sys, getopt
import os, errno
import subprocess
import shutil
import time
from collections import OrderedDict
import logging as log

starttime = runtime = time.time()
log.basicConfig(level=log.INFO)

def ellapsed(overall=False):
    global runtime
    t = (time.time() - starttime) if overall else (time.time() - runtime)
    runtime = time.time()
    return "%d:%02d:%02d" % (int(t//3600), int(t/60)%3600, round(t)%60)

# Mkdir (and erase if exist)
def makedir(path):
    shutil.rmtree(path, True)
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path): pass
        else: raise

# Get Size of a folder (recursively)
def getFolderSize(folder):
    total_size = os.path.getsize(folder)
    for item in os.listdir(folder):
        itempath = os.path.join(folder, item)
        if os.path.isfile(itempath):
            total_size += os.path.getsize(itempath)
        elif os.path.isdir(itempath):
            total_size += getFolderSize(itempath)
    return total_size

# Utils to calculate the ratio between
#   the current execution time and the movie duration
#   gives movie_duration * factor = encoding_duration
def time_factor(filepath):
    factor = 0
    def result(valid=True):
        if valid:
            log.info('Time factor: x{0}'.format(round(factor, 2)))
        return factor

    # FFProbe command
    ffprobecmd = 'ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1 {0}'.format(filepath)
    probe = subprocess.Popen(ffprobecmd.split(), stdout=subprocess.PIPE)
    output, err = probe.communicate()
    try:
        output = float(output.split('=')[1])
        factor = (time.time() - starttime)/output
        return result()
    except:
        pass
    return result(False)

# Utils to calculate the ratio between
#   the output folder size and the original file size
#   gives original_size * factor = output_size
def size_factor(filepath):
    factor = 0
    def result(valid=True):
        if valid:
            log.info('Size factor: x{0}'.format(round(factor, 2)))
        return factor

    try:
        inputsize = os.path.getsize(filepath)
        outputsize = getFolderSize(os.path.splitext(filepath)[0])
        factor = outputsize*1./inputsize
        return result()
    except:
        pass
    return result(False)

# Utils to detect video aspect ratio using FFProbe
def detect_ratio(filepath):
    # detect video aspect ratio
    aspectratio = '16/9'
    width = height = 0

    def result(valid=True):
        if valid:
            log.info('Detected aspect ratio: {0}'.format(aspectratio))
        else:
            log.warning('Can\'t detect video aspect ratio.. going with default: {0}'.format(aspectratio))
        return aspectratio

    # FFProbe command
    ffprobecmd = 'ffprobe -v error -show_entries stream=width,height -of default=noprint_wrappers=1 {0}'.format(filepath)
    probe = subprocess.Popen(ffprobecmd.split(), stdout=subprocess.PIPE)
    output, err = probe.communicate()
    output = output.split()
    try:
        width = float(output[ (0 if output[0][0] == 'w' else 1) ].split('=')[1])
        height = float(output[ (0 if output[0][0] == 'h' else 1) ].split('=')[1])
        ratio = width/height
        if ratio > 1.2 and ratio < 1.9:
            if ratio < 1.5: aspectratio = '4/3'
            else: aspectratio = '16/9'
            return result()
    except:
        pass
    return result(False)


# Utils to detect max segment bitrate using FFProbe
def detect_bitrate(path):
    bitrate = 0

    # FFProbe command
    ffprobecmd = 'ffprobe -v error -show_entries format=bit_rate -of default=noprint_wrappers=1 {0}'

    # Check each file in the path and keep max bitrate
    for (dirpath, dirnames, files) in os.walk(path):
        for filename in files:
            if filename[-2:] == 'ts':
                try:
                    filepath = os.path.join(path, filename)
                    probe = subprocess.Popen(ffprobecmd.format(filepath).split(), stdout=subprocess.PIPE)
                    output, err = probe.communicate()
                    output = int(output.split('=')[1])
                    if not output > 0: raise
                    bitrate = max(bitrate, output)
                except:
                    log.error('Can\'t detect bitrate for the file {0}'.format(filepath))

    log.debug('Bitrate detected: {0}'.format(bitrate)) if bitrate > 0 else  log.warning('No Bitrate detected in {0}'.format(path))
    return bitrate

# Force Key Frames to match segment segmentsize
def setkeyframes(inputfile, interval=10):
    # base ffmpeg command
    ffmpegcmd = 'ffmpeg -i {input} -force_key_frames "expr:gte(t,n_forced*2)" k{interval}-{input}'.format(input=inputfile, interval=interval)
    print ffmpegcmd
    log.info('Forcing Keyframes..')

    # Execute FFMPEG command
    probe = subprocess.Popen(ffmpegcmd, stdout=subprocess.PIPE, shell=True)
    output, err = probe.communicate()

    if not err:
        log.info('Forced Keyframe media created: k{interval}-{input}'.format(input=inputfile, interval=interval))
        return 'k{interval}-{input}'.format(input=inputfile, interval=interval)
    else:
        log.warning('Can\'t force keyframes..')
        return inputfile


# Execute FFMPEG segmenter based on predefined profiles
def segmenter(inputfile, dryrun=False, urlprefix='../'):
    # import profiles presets
    import presets

    # base ffmpeg command for segmenter
    ffmpegcmd_segm = 'ffmpeg -y -v error -i %(inputfile)s -c:a aac -strict experimental -ac 2 -b:a %(audiobitrate)s -ar 44100 -c:v libx264 -pix_fmt yuv420p -profile:v %(profile)s -preset %(ffmpegmode)s -level %(level)s -b:v %(videobitrate)s -maxrate %(videobitrate)s -bufsize %(buffersize)s -threads 0 -r %(fps)s -g %(gop)s -hls_time %(segmentsize)s -hls_list_size 0 -s %(resolution)s %(outputpath)s/%(outputname)s.m3u8'

    # fallback media
    ffmpegcmd_media = 'ffmpeg -y -v error -i %(inputfile)s -c:a aac -strict experimental -ac 2 -b:a %(audiobitrate)s -ar 44100 -c:v libx264 -pix_fmt yuv420p -profile:v %(profile)s -preset %(ffmpegmode)s -level %(level)s -b:v %(videobitrate)s -maxrate %(videobitrate)s -bufsize %(buffersize)s -threads 0 -r %(fps)s -s %(resolution)s %(workpath)s/%(workname)s.mp4'


    # working path / name
    workname = os.path.splitext(os.path.basename(inputfile))[0]
    workpath = os.path.splitext(inputfile)[0]
    if not dryrun:
        makedir(workpath)


    # prepare ffmpeg commands formated with the profiles
    variants = OrderedDict()
    profiles = presets.build(inputfile=inputfile, ratio=detect_ratio(inputfile))
    ellapsed()

    # try to force keyframes
    # inputfile = setkeyframes(inputfile=inputfile, interval=presets.SEGMENT_SIZE)
    ellapsed()


    # start segmenter
    firstPreset = True
    for quality, preset in profiles:
        preset['outputpath'] = os.path.join(workpath, quality)
        preset['workname'] = workname
        preset['workpath'] = workpath
        err = False
        if not dryrun:
            # Output path
            makedir(preset['outputpath'])

            # Execute FFMPEG command MEDIA (for the first preset only)
            if firstPreset:
                log.info("Encoding fallback media with profile: {0}".format(quality))
                probe = subprocess.Popen((ffmpegcmd_media % preset).split(), stdout=subprocess.PIPE)
                output, err = probe.communicate()
                firstPreset = False

            # Execute FFMPEG command SEGMENTER
            log.info("Processing media with profile: {0}".format(quality))
            probe = subprocess.Popen((ffmpegcmd_segm % preset).split(), stdout=subprocess.PIPE)
            output, err = probe.communicate()

        # Save m3u8 file path, resolution and maximum bitrate
        if not err:
            variants[quality] = {
                'resolution': preset['resolution'],
                'playlist': urlprefix + os.path.join(workname, quality, preset['outputname']+'.m3u8'),
                'bitrate': detect_bitrate(preset['outputpath'])
            }
        # Log
        try:
            if variants[quality]['bitrate'] > 0:
                log.debug('Profile {0} completed in {1}'.format(quality, ellapsed()))
        except:
            log.error('Error with profile {0}: {1}'.format(quality, err))

    # Create variants playlist
    variantfilepath = os.path.join(workpath, workname+'.m3u8')
    with open(variantfilepath, 'w') as f:
        f.write('#EXTM3U\n')
        for quality, variant in variants.items():
            if variant['bitrate'] > 0:
                f.write('#EXT-X-STREAM-INF:BANDWIDTH=%(bitrate)s,RESOLUTION=%(resolution)s\n' % variant)
                f.write(variant['playlist']+'\n')
        log.info('Variant playlist created: {0}'.format(variantfilepath))



def main(argv):
    # Init
    usage = 'Usage: {0} -i <input file> [-u <urlprefix>] [-p] '.format(sys.argv[0])
    inputfile = test = variantonly = False
    urlprefix = '../'
    # Parse Args
    try:
        opts, args = getopt.getopt(argv,"hi:u:p",[])
    except getopt.GetoptError:
        print usage
        sys.exit(2)
    for opt, arg in opts:
        if opt in ("-i"):
            inputfile = arg
        if opt in ("-u"):
            urlprefix = arg
        if opt in ("-p"):
            variantonly = True

    if not inputfile:
        print usage
        sys.exit(0)

    # check if input file exists
    if not os.path.isfile(inputfile):
        log.error('File "{0}" not found'.format(inputfile))
        sys.exit(2)

    # convert files and create playlists
    segmenter(inputfile=inputfile, dryrun=variantonly, urlprefix=urlprefix)

    # Exit
    log.info('HLS segmentation completed in {0}'.format( ellapsed(True) ))
    time_factor(inputfile)
    size_factor(inputfile)
    sys.exit(0)


if __name__ == "__main__":
   main(sys.argv[1:])


#print 'Number of arguments:', len(sys.argv), 'arguments.'
#print 'Argument List:', str(sys.argv)

# lowprofile = 'ffmpeg -y -i {0} -c:a aac -strict experimental -ac 2 -b:a 64k -ar 44100 -c:v libx264 -pix_fmt yuv420p -profile:v baseline -level 21 -b:v 400K -r 12 -g 36 -f hls -hls_time 1 -hls_list_size 0 -s 480x270 {1}_480x270.m3u8'
# medprofile = 'ffmpeg -y -i {0} -c:a aac -strict experimental -ac 2 -b:a 96k -ar 44100 -c:v libx264 -pix_fmt yuv420p -profile:v baseline -level 31 -b:v 600K -r 24 -g 72 -f hls -hls_time 1 -hls_list_size 0 -s 640x360 {1}_640x360.m3u8'
# hiprofile = 'ffmpeg -y -i {0} -c:a aac -strict experimental -ac 2 -b:a 96k -ar 44100 -c:v libx264 -pix_fmt yuv420p -profile:v main -level 32 -b:v 1500K -r 24 -g 72 -f hls -hls_time 1 -hls_list_size 0 -s 1280x720 {1}_1280x720.m3u8'


#ffmpeg -i sintel_trailer-1080p.mp4 -hls_time 1 -hls_list_size 0 -hls_allow_cache 1 -hls_base_url 'http://localhost/' -strict -2 out.m3u8