#!/usr/bin/env bash
# -------------------------------------------------------
#  Video Stabilizer based on MLT and vid.stab
#  For installation instructions, please check
#  http://bernaerts.dyndns.org/linux/350-ubuntu-xenial-rotate-stabilize-video-melt-vidstab
#  
#  Depends on :
#    * melt (from ppa:sunab/kdenlive-release)
#    * vidstab (from ppa:sunab/kdenlive-release)
#    * yad
#    * exiftool
#    * avconv and x264
#
#  Revision history :
#    04/03/2012, V1.0 - Creation by N. Bernaerts
#    04/04/2012, V1.1 - Add codecs as parameters
#    17/01/2015, V2.0 - Completle rewrite for Ubuntu 14.04 LTS 
#                       Handle rotation as well thanks to Guy Eagling
#    19/01/2015, V2.1 - Manage per process .trf file to allow parallel execution
#    19/01/2015, V2.2 - Add .pid file to handle processing cancellation
#    23/01/2015, V3.0 - Check tools availability, 
#                       Detect available stabilization filters (vidstab, videostab2 or videostab)
#                       Add Fedora compatibility (mlt-melt) thanks to Guy Eagling
#    06/03/2015, V3.1 - Externalize parameters to ~/.config/video-stabilize.conf
#                       Add multiple files processing thanks to Hingo's idea
#    10/03/2015, V3.2 - Update exiftool handling for MKV files
#    12/03/2015, V3.3 - Add extra parameters to dialog box 
#                       Change progress calculation to avoid errors
#    29/05/2015, V3.4 - Check presence of ~/.config/video-stabilize.conf 
#    26/07/2015, V3.5 - Change options of first pass to avoid some errors (thanks to Gustavo Lapido Loureiro)
#                       Add number of thread as encoder option
#                       Kill process and parent process in case of cancellation
#    22/09/2015, V3.6 - Correction of a bug in final encoding options (thanks to Ted Bartlett)
#    21/02/2016, V3.7 - Handle scientific notation for video rate (thanks to Aslanex)
#    24/08/2016, V3.8 - Use specific build of melt for vid.stab (Ubuntu Xenial 16.04)
# -------------------------------------------------------

# -------------------------------------------------
#          Check tools availability
# -------------------------------------------------

# check /opt/mlt-vidstab/melt (specific build with vid.stab filter) or melt (ubuntu) or mlt-melt (fedora)
[ -f /opt/vidstab/melt ] && CMD_MELT="/opt/vidstab/melt" || CMD_MELT="melt"
command -v ${CMD_MELT} >/dev/null 2>&1 || { CMD_MELT="mlt-melt"; }
command -v ${CMD_MELT} >/dev/null 2>&1 || { zenity --error --text="[error] Please install Melt from MLT framework"; exit 1; }

# check yad, exiftool and x264
command -v yad >/dev/null 2>&1 || { zenity --error --text="[error] Please install Yad"; exit 1; }
command -v exiftool >/dev/null 2>&1 || { zenity --error --text="[error] Please install ExifTool"; exit 1; }
command -v x264 >/dev/null 2>&1 || { zenity --error --text="[error] Please install x264"; exit 1; }

# -------------------------------------------------------
#      Read parameters from configuration file
# -------------------------------------------------------

# Configuration file : ~/.config/video-stabilize.conf
FILE_CONF="$HOME/.config/video-stabilize.conf"

# check configuration file
[ -f "${FILE_CONF}" ] || { zenity --error --text="[error] Please create and configure ${FILE_CONF}"; exit 1; }

# Load configuration file
ARR_EXT=( $(cat "${FILE_CONF}" | grep "extension" | cut -d'=' -f 2- | cut -d ";" --output-delimiter=" " -f 1-) )
STAB_SHAKINESS=$(cat "${FILE_CONF}" | grep "shakiness" | cut -d'=' -f2)
STAB_SMOOTHING=$(cat "${FILE_CONF}" | grep "smoothing" | cut -d'=' -f2)
STAB_OPTZOOM=$(cat "${FILE_CONF}" | grep "optzoom" | cut -d'=' -f2)
STAB_EXTRA=$(cat "${FILE_CONF}" | grep "extra" | cut -d'=' -f2-)
VIDEO_CODEC=$(cat "${FILE_CONF}" | grep "video-codec" | cut -d'=' -f2)
AUDIO_CODEC=$(cat "${FILE_CONF}" | grep "audio-codec" | cut -d'=' -f2)
VIDEO_RATE=$(cat "${FILE_CONF}" | grep "video-rate" | cut -d'=' -f2)
AUDIO_RATE=$(cat "${FILE_CONF}" | grep "audio-rate" | cut -d'=' -f2)
ENCODE_OPTION=$(cat "${FILE_CONF}" | grep "option" | cut -d'=' -f2-)
ENCODE_THREAD=$(cat "${FILE_CONF}" | grep "thread" | cut -d'=' -f2-)

# -------------------------------------------------------
#          Retrieve or select video file
# -------------------------------------------------------

IFS=$'\n'

# loop thru parameters to check for video extensions
while test ${#} -gt 0
do
  # if parameter is a file
  if [ -f "$1" ]
  then
    # get file extension
    FILE_EXT=$(echo "$1" | sed 's/^.*\.\(.*\)$/\1/')
    
    # check if file extension given in parameter is in the allowed extension list
    EXT_OK=$(echo "${ARR_EXT[@]}" | grep --ignore-case "${FILE_EXT}")

    # if ok, add it to the video files array
    [ "${EXT_OK}" != "" ] && ARR_VIDEO=( "${ARR_VIDEO[@]}" "$1" )
  fi
  
  # go to next parameter
  shift
done

# if there is no candidate files
if [ ${#ARR_VIDEO[@]} -eq 0 ]
then
  # generate allowed extension list
  LIST_EXT=$(echo "*.${ARR_EXT[@]}" | sed 's/ / *\./g')

  # open multiple files selection dialog box
  LST_VIDEO=$(yad --center --width=800 --height=500 --window-icon "video" --image "stabilizer" \
                  --file --multiple --file-filter="Video file (${LIST_EXT[@]})|${LIST_EXT[@]}" \
                  --title="Select video file to stabilize")

  # generate video files array
  ARR_VIDEO=($(echo "${LST_VIDEO}" | tr "|" "\n"))
fi 

# ---------------------------------------------------------
#   Analyse video files and select processing parameters
# ---------------------------------------------------------

# loop thru selected video files
for FILE_PATH in "${ARR_VIDEO[@]}"
do
  # generate temporary exif file
  FILE_EXIF=$(mktemp -t "stab-XXXXXXXX.exif")

  # get video metadata
  exiftool "${FILE_PATH}" > "${FILE_EXIF}"
  FILE_WIDTH=$(cat "${FILE_EXIF}" | grep "^Image Width" | cut -d':' -f2 | xargs)
  FILE_HEIGHT=$(cat "${FILE_EXIF}" | grep "^Image Height" | cut -d':' -f2 | xargs)
  FILE_BITRATE=$(cat "${FILE_EXIF}" | grep "^Avg Bitrate" | cut -d':' -f2 | xargs)
  FILE_ROTATE=$(cat "${FILE_EXIF}" | grep "^Rotation" | cut -d':' -f2 | xargs)
  [ "${FILE_BITRATE}" == "" ] && FILE_BITRATE="Unknown"
  [ "${FILE_ROTATE}" == "" ] && FILE_ROTATE="0"

  # remove temporary exif file
  rm "${FILE_EXIF}"

  # get encoding parameters using local number decimal separator (, or .)
  SEPARATOR=$(printf "%'.2f" 1 | sed 's/^1\(.\).*$/\1/')
  VIDEO_RATIO=$(echo "scale=2; ${FILE_WIDTH} / ${FILE_HEIGHT}" | bc | sed 's/[\.\,]/'${SEPARATOR}'/g')
  [ "${FILE_ROTATE}" = "0" -o "${FILE_ROTATE}" = "180" ] && LST_RATIO="1|${VIDEO_RATIO}" || LST_RATIO="${VIDEO_RATIO}|1" 

  # detect if rotation needed
  [ "${FILE_ROTATE}" = "0" ] && CHECK_ROTATE="FALSE" || CHECK_ROTATE="TRUE"

  # set title and text of dialog box
  TITLE="${FILE_PATH}"
  TEXT="Select transformation parameters :\n ( size = ${FILE_WIDTH}x${FILE_HEIGHT}, rate = ${FILE_BITRATE} )\n"

  # get list of stabilization filters available from current melt version
  ARR_FILTER=$(${CMD_MELT} -query filters | grep stab | awk '{print $2}' | sort -r)
  LST_FILTER=$(echo "${ARR_FILTER}" | sed 's/ /|/g')

  # display dialog box
  CHOICE=$(yad --title "${TITLE}" --text "${TEXT}" --center --window-icon "video" --image "stabilizer" --width 500 \
    --form --item-separator='|' \
    --field="Rotate:CHK" "${CHECK_ROTATE}" \
    --field="  - Angle (${FILE_ROTATE}° detected):NUM" "${FILE_ROTATE}|0..359" \
    --field="  - Resize ratio (${VIDEO_RATIO} for 90°):CB" "${LST_RATIO}" \
    --field="Stabilize:CHK" "TRUE" \
    --field="  - Stabilization filter:CB" "${LST_FILTER}" \
    --field="  - Shakiness [ 0 ... 10 ]:NUM" "${STAB_SHAKINESS}|1..10" \
    --field="  - Smoothing [ 0 ... 100 ]:NUM" "${STAB_SMOOTHING}|0..100"  \
    --field="  - Optimal Zoom [ 0, 1, 2 ]:NUM" "${STAB_OPTZOOM}|0..2" \
    --field="  - Extra stabilize options" "${STAB_EXTRA}" \
    --field="Encoder:LBL" "final" \
    --field="  - ${VIDEO_CODEC} video bitrate (Kbits/s):NUM" "${VIDEO_RATE}|1..50000|100" \
    --field="  - ${AUDIO_CODEC} audio bitrate (Kbits/s):NUM" "${AUDIO_RATE}|1..320" \
    --field="  - Extra encoder parameters" "${ENCODE_OPTION}" )

  # retrieve parameters
  ROTATE=$(echo "${CHOICE}" | cut -d'|' -f1)
  ROTATE_ANGLE=$(echo "${CHOICE}" | cut -d'|' -f2)
  ROTATE_RATIO=$(echo "${CHOICE}" | cut -d'|' -f3)
  STABILIZE=$(echo "${CHOICE}" | cut -d'|' -f4)
  STAB_FILTER=$(echo "${CHOICE}" | cut -d'|' -f5)
  STAB_SHAKINESS=$(echo "${CHOICE}" | cut -d'|' -f6)
  STAB_SMOOTHING=$(echo "${CHOICE}" | cut -d'|' -f7)
  STAB_OPTZOOM=$(echo "${CHOICE}" | cut -d'|' -f8)
  STAB_EXTRA=$(echo "${CHOICE}" | cut -d'|' -f9)
  VIDEO_RATE=$(echo "${CHOICE}" | cut -d'|' -f11 | tr ',' '.' | awk '{ print sprintf("%.0f", $1); }')
  AUDIO_RATE=$(echo "${CHOICE}" | cut -d'|' -f12)
  ENCODE_OPTION=$(echo "${CHOICE}" | cut -d'|' -f13)

  # if it is needed, add file to processing queue
  if [ "${STABILIZE}" = "TRUE" ] || [ "${ROTATE}" = "TRUE" ];
  then
    ARR_FILE=("${ARR_FILE[@]}" "${FILE_PATH}")
    ARR_ROTATE=("${ARR_ROTATE[@]}" "${ROTATE}")
    ARR_ROTATE_ANGLE=("${ARR_ROTATE_ANGLE[@]}" "${ROTATE_ANGLE}")
    ARR_ROTATE_RATIO=("${ARR_ROTATE_RATIO[@]}" "${ROTATE_RATIO}")
    ARR_STABILIZE=("${ARR_STABILIZE[@]}" "${STABILIZE}")
    ARR_STAB_FILTER=("${ARR_STAB_FILTER[@]}" "${STAB_FILTER}")
    ARR_STAB_SHAKINESS=("${ARR_STAB_SHAKINESS[@]}" "${STAB_SHAKINESS}")
    ARR_STAB_SMOOTHING=("${ARR_STAB_SMOOTHING[@]}" "${STAB_SMOOTHING}")
    ARR_STAB_OPTZOOM=("${ARR_STAB_OPTZOOM[@]}" "${STAB_OPTZOOM}")
    ARR_STAB_EXTRA=("${ARR_STAB_EXTRA[@]}" "${STAB_EXTRA}")
    ARR_VIDEO_RATE=("${ARR_VIDEO_RATE[@]}" "${VIDEO_RATE}")
    ARR_AUDIO_RATE=("${ARR_AUDIO_RATE[@]}" "${AUDIO_RATE}")
    ARR_ENCODE_OPTION=("${ARR_ENCODE_OPTION[@]}" "${ENCODE_OPTION}")
  fi
done

# -------------------------------------------------------
#     Process files for rotation and stabilization 
# -------------------------------------------------------

# loop thru video files to process
NUM_FILE=${#ARR_FILE[@]}
for ((INDEX=0; INDEX < NUM_FILE; INDEX++))
do
  # --------------------------------
  #         Get file data
  # --------------------------------
  # get current file path 
  FILE_PATH="${ARR_FILE[$INDEX]}"

  # generate the filenames
  FILE_BASE=$(echo "${FILE_PATH}" | sed "s|^\(.*\)\..*$|\1|")
  FILE_STAB="${FILE_BASE}-stab.mp4"

  # generate temporary files
  DIR_TMP=$(mktemp -d "$HOME/.stab-XXXXXXXX")
  FILE_TRF="${DIR_TMP}/video.trf"
  FILE_PID="${DIR_TMP}/video.pid"
  FILE_MLT="${DIR_TMP}/video.mlt"
  FILE_LOG="${DIR_TMP}/video.log"

  # get parameters
  ROTATE="${ARR_ROTATE[$INDEX]}"
  ROTATE_ANGLE="${ARR_ROTATE_ANGLE[$INDEX]}"
  ROTATE_RATIO="${ARR_ROTATE_RATIO[$INDEX]}"
  STABILIZE="${ARR_STABILIZE[$INDEX]}"
  STAB_FILTER="${ARR_STAB_FILTER[$INDEX]}"
  STAB_SHAKINESS="${ARR_STAB_SHAKINESS[$INDEX]}"
  STAB_SMOOTHING="${ARR_STAB_SMOOTHING[$INDEX]}"
  STAB_OPTZOOM="${ARR_STAB_OPTZOOM[$INDEX]}"
  STAB_EXTRA="${ARR_STAB_EXTRA[$INDEX]}"
  VIDEO_RATE="${ARR_VIDEO_RATE[$INDEX]}"
  AUDIO_RATE="${ARR_AUDIO_RATE[$INDEX]}"
  ENCODE_OPTION="${ARR_ENCODE_OPTION[$INDEX]}"

  (
  # initilize transformation arrays
  unset ARR_STABILIZE
  unset ARR_ROTATE
  unset ARR_ENCODE

  # ------------------------------------------------
  #    PREPARATION : Setup processing parameters
  # ------------------------------------------------

  # initial display
  echo "# Computing parameters"
  echo "0"

  # if needed, generate stabilizer parameters array
  if [ "${STABILIZE}" = "TRUE" ]
  then
    # get default parameters
    ARR_STABILIZE=( $(echo "${STAB_EXTRA}" | cut -d ";" --output-delimiter=" " -f 1-) )
    
    # add current parameters
    ARR_STABILIZE=( "-filter" "${STAB_FILTER}" "filename=${FILE_TRF}" "shakiness=${STAB_SHAKINESS}" "smoothing=${STAB_SMOOTHING}" "optzoom=${STAB_OPTZOOM}" "${ARR_STABILIZE[@]}" )
  fi
  
  # if needed, set rotation filter
  if [ "${ROTATE}" = "TRUE" ]
  then
    # convert angle for filter use
    X_ANGLE=$((10 * ${ROTATE_ANGLE} / 2))

    # set rotation filter
    ARR_ROTATE=( "-filter" "affine" "transition.fix_rotate_x=${X_ANGLE}" )

    # if different than 1, set transformation ratio
    [ "${ROTATE_RATIO}" != "1" ] && ARR_ROTATE=( "${ARR_ROTATE[@]}" "transition.scale_x=${ROTATE_RATIO}" "transition.scale_y=${ROTATE_RATIO}" )
  fi

  # ----------------------------------------------
  #   PASS 1 : Rotation / Stabilization analysis
  # ----------------------------------------------

  # information display
  echo "# Analysing file"

  # launch analysis command in background and get its process id
  ${CMD_MELT} -progress "${FILE_PATH}" "${ARR_ROTATE[@]}" "${ARR_STABILIZE[@]}" -consumer "xml:${FILE_MLT}" all=1 2>"${FILE_LOG}" &
  PID=$!

  # save current shell and background process id
  echo "$BASHPID" > "${FILE_PID}"
  echo "${PID}" >> "${FILE_PID}"

  while [ "${PID}" != "" ] 
  do    
    # check if process is still running
    PID=$(ps aux | awk '{print $2}' | grep "${PID}")

    # calculate process completion
    LOG_LINE=$(tail --bytes=50 "${FILE_LOG}" | grep "Frame" | grep "percentage")
    LOG_FRAME=$(echo "${LOG_LINE}" | sed 's/^.*Frame[ :]*\([0-9]*\).*$/\1/')
    LOG_PERCENT=$(echo "${LOG_LINE}" | sed 's/^.*percentage[ :]*\([0-9]*\).*$/\1/')

    # if percentage has been retrieved from log, display process completion
    [[ ${LOG_FRAME} == +([0-9]) ]] && echo "# Analysing frame ${LOG_FRAME}"
    [[ ${LOG_PERCENT} == +([0-9]) ]] && echo $((${LOG_PERCENT} / 2))

    # wait for 2 seconds
    sleep 2
  done

  # ------------------------------------------
  #   PASS 2 : Final result encoding
  # ------------------------------------------

  # information display
  echo "# Generating file"

  # get default encoding parameters
  ARR_ENCODE=( $(echo "${ENCODE_OPTION}" | cut -d ";" --output-delimiter=" " -f 1-) )
  
  # add current encoding parameters
  ARR_ENCODE=( "vcodec=${VIDEO_CODEC}" "b=${VIDEO_RATE}k" "acodec=${AUDIO_CODEC}" "ab=${AUDIO_RATE}k" "${ARR_ENCODE[@]}" )

  # launch generation command in background and get its process id
  ${CMD_MELT} -progress "${FILE_MLT}" -audio-track "${FILE_PATH}" -consumer "avformat:${FILE_STAB}" "${ARR_ENCODE[@]}" real_time=-${ENCODE_THREAD} 2>"${FILE_LOG}" &
  PID=$!

  # save current shell and background process id
  echo "$BASHPID" > "${FILE_PID}"
  echo "${PID}" >> "${FILE_PID}"

  # follow the stabilization progress
  while [ "${PID}" != "" ] 
  do    
    # check if process is still running
    PID=$(ps aux | awk '{print $2}' | grep "${PID}")

    # calculate process completion
    LOG_LINE=$(tail --bytes=50 "${FILE_LOG}" | grep "Frame" | grep "percentage")
    LOG_FRAME=$(echo "${LOG_LINE}" | sed 's/^.*Frame[ :]*\([0-9]*\).*$/\1/')
    LOG_PERCENT=$(echo "${LOG_LINE}" | sed 's/^.*percentage[ :]*\([0-9]*\).*$/\1/')

    # if percentage has been retrieved from log, display process completion
    [[ ${LOG_FRAME} == +([0-9]) ]] && echo "# Generating frame ${LOG_FRAME}"
    [[ ${LOG_PERCENT} == +([0-9]) ]] && echo $((50 + (${LOG_PERCENT} / 2)))

    # wait for 2 seconds
    sleep 2
  done

  # remove PID file
  rm "${FILE_PID}"

  # declare end of processing
  echo "# File processed and available"
  echo "100"
  ) | yad --center --width=600 --window-icon "video" --image "stabilizer" --progress --auto-close --title "[$((${INDEX} + 1))/${NUM_FILE}] ${FILE_PATH}"

  # -----------------------------------------
  #   END : Final process and files cleanup
  # -----------------------------------------

  # if process is still running (operation has been canceled), kill the process
  [ -f "${FILE_PID}" ] && kill -9 $(cat "${FILE_PID}")

  #  cleaning-up of all temporary files
  rm -R ${DIR_TMP}

done