vidstab.py

#!/usr/bin/env python3

"""
vidstab - Performs video stabilization on input using vid.stab in ffmpeg

Usage:
  vidstab <input> [--with-trf=TRFILE | --transform-only] [options]
  vidstab (-h | --help)
  vidstab --version

Options:
  -s --shakiness=NUM    Shakiness level for vidstabdetect. Values range from 1
                        to 10. [default: 5]
  -r --result=PATH      Location for transformation file from vidstabdetect.
  -a --accuracy=NUM     Accuracy level for vidstabdetect. Values range from 1 to
                        15. [default: 15]
  -S --stepsize=NUM     Stepsize for vidstabdetect motion search. [default: 6]
  --mincontrast=NUM     A floating point value from 0-1. Discard measurement
                        fields with contrast below this value. [default: 0.3]
  -3 --tripod=FRAME     Reference frame number for the simulation of static
                        viewpoint, Disabled if left at 0. [default: 0]
  --show=NUM            Show fields and transforms with 1 or 2. [default: 0]
  -x --transform-only   Calculate the TR File and then quit.
  -t --with-trf=TRFILE  Set a TR File that has already been calculated.
  -m --smoothing=NUM    Number of frames to use for lowpass filtering of camera
                        movements. Greater is smoother but slower movements. 0
                        is a special case for static camera, --tripod will
                        override this option if used. [default: 10]
  -o --optalgo=ALGO     Set the camera path optimization algorithm. Values are:
                        "gauss" and "avg": [default: gauss]
  -d --maxshift=NUM     Set maximum number of pixels allowed to translate
                        frames. -1 is unlimited. [default: -1]
  -a --maxangle=RAD     Set maximum angle in radians allowed to rotate frames.
                        Value of -1 is unlimited. [default: -1]
  -C --crop=METHOD      How to deal with empty borders after frame movement
                        compensation. "keep" will use info from prior frame and
                        "black" will do nothing (black-fill). [default: keep]
  -i --invert           Invert transforms if set.
  -k --relative         Make transforms relative to previous frame (instead of
                        (absolute) if set. --tripod will override this.
  -z --zoom=NUM         Set percentage to zoom, positive will zoom in while
                        negative will zoom out. [default: 0]
  -Z --optzoom=NUM      Calculate optimal zooming to avoid blank borders. 0 is
                        disabled, 1 is optimal static zoom, 2 is optimal
                        adaptive zoom. [default: 1]
  --zoomspeed=NUM       Maximum zoom percentage allowed each frame for optimal
                        adaptive zoom. 0 to 5. [default: 0.25]
  -I --interpol=METHOD  Set interpolation method. One of: no, linear, bilinear
                        bicubic (slow calculation). [default: bilinear]
  -c --crf=NUM          libx264 CRF value to use for output encoding [default: 18]
  -p --preset=PRESET    libx264 preset to use [default: medium]
  -t --tune <tune>      libx264 tune setting to use [default: film]
  -h --help             Show this screen.
  --version             Show version.

https://github.com/georgmartius/vid.stab
"""

from docopt import docopt
import subprocess
import os.path
import sys

if __name__ == '__main__':
    args = docopt(__doc__, version='vidstab 1.0')
    print(args)

    #sys.exit()
    
    basename, _ext = os.path.splitext(args['<input>'])
    
    if args['--with-trf'] is None:
        #set the options up as well as the transform file name
        if args['--result'] is None:
            args['--result'] = basename + '.trf'
        subprocess.run(['ffmpeg', '-i', args['<input>'],
                        '-vf', '''vidstabdetect=stepsize={--stepsize}:\
shakiness={--shakiness}:accuracy={--accuracy}:tripod={--tripod}:\
result={--result}:mincontrast={--mincontrast}'''.format(**args),
                        '-f', 'null', '-'])
    
    if not args['--transform-only']:
        if args['--with-trf'] is None:
            args['--with-trf'] = basename + '.trf'
        if args['--invert']:
            args['--invert'] = '1'
        else:
            args['--invert'] = '0'
        if args['--relative']:
            args['--relative'] = '1'
        else:
            args['--relative'] = '0'
        if args['--tripod'] != '0':
            print('tripod mode setting --relative and --smoothing to 0')
            args['--relative'] = 0
            args['--smoothing'] = 0
            
        command = ['ffmpeg', '-i', args['<input>'],
                   '-map', '0',
                   '-vf', '''vidstabtransform=input={--with-trf}:\
smoothing={--smoothing}:optalgo={--optalgo}:maxshift={--maxshift}:\
maxangle={--maxangle}:crop={--crop}:invert={--invert}:relative={--relative}:\
zoom={--zoom}:optzoom={--optzoom}:zoomspeed={--zoomspeed}:\
interpol={--interpol},unsharp=5:5:0.8:3:3:0.4'''.format(**args),
                   '-c:v', 'libx264',
                   '-preset', args['--preset'],
                   '-tune',   args['--tune'],
                   '-crf',    args['--crf'],
                   '-c:a', 'copy', 'vidstabbed_' + basename + '.mkv'
                   ]
        print(' '.join(command))
        subprocess.run(command)