stackalign

#! /bin/bash
# align - align image stack with focus-stack

usage() {
  echo "
align - align image stack

Usage:
  align [OPTIONS] IMAGES ...

Options:
 --help                    Show this help.
 --output DIR              Store aligned images in directory DIR.
                           Default: ./aligned
 --rotate DEGREE           Rotate images clockwise by degree DEGREE.
 --contrast [=yes|no]      Correct contrast and exposure differences.
 --whitebalance [=yes|no]  Correct white balance differences.
 -V, --showimage           Show aligned images in image viewer geeqie.

Depends on:
 - focus-stack from Petteri Aimonen:
   https://github.com/PetteriAimonen/focus-stack
 - ImageMagick
 - geeqie (optional)

NOTE: You might need to remove the orientation tag of the source images first.
  Possible tool and command: 
  exiftool -n -Orientation=0 -overwrite_original IMAGES
NOTE: If the images are given in a lossy format like JPG, you should convert
  them first to a lossless format like TIF or PNG. stackprepare can do that.
"
}
error() {
  echo "stackalign ERROR: $*" >&2
  finish 1
}

### misc
digitonly() {
  echo "${1//[^0-9.]/}"
}
printsameline() {
  # print $1 without newline at begin of current line
  echo -ne "${1:-}\033[0K\r" >&2
}
showimage() {
  [ "$Showimageprocessing" = "yes" ] && geeqie -t -r --File:"${1:-}" 2>/dev/null &
  return 0
}

### run commands on all CPUs
multicore() {
  # Run multiple processes in parallel, but not more than $Multicore_maxprocesses
  # $1 Command
  # $2 Image to show if $1 is finished
  # Run multicore_wait afterwards to wait for the last processes to finish.

  local Process Command
  local Mem_needed Zram

  [ "${1:-}" = "-t1" ] && {
    shift
    [ "$Multicore_processcount" -gt 0 ] && {
      multicore_wait || return 1
    }
  }

  [ "$Multicore_processcount" = "$Multicore_maxprocesses" ] && {
    multicore_wait || return 1
  }
  [ "$Multicore_processcount" = "0" ] && Multicore_memorymax="$(printfreememory)"  
  Mem_needed=0
  for Process in $(seq $Multicore_maxprocesses); do
    Mem_needed="$(awk "BEGIN {print $Mem_needed + ${Multicore_memory[$Process]:-0} }" )"
  done
  Mem_needed="$((Mem_needed+${3:-0}))"
  [ "$Mem_needed" -gt "$Multicore_memorymax" ] && {
    note "multicore: Low memory. Waiting for $Multicore_processcount running processes to finish. Need: $((Mem_needed/1000)) MB, Available: $((Multicore_memorymax/1000)) MB"
    [ "$Mem_needed" -gt "$Multicore_memorymax" ] && [ "$Multicore_processcount" = "0" ] && note "multicore: Likely hard disk cache will be used and slow down the calculation."
    multicore_wait || return 1
  }

  ifcmdbreak && return 1

  Multicore_processcount=$((Multicore_processcount +1))

  Command="$(cut -d ' ' -f1 <<< "${1:-}")"
  case $(type -t "$Command") in
    file) Command="nice ${1:-}" ;;
    *)    Command="${1:-}" ;;
  esac

  #verbose "multicore: ${1:-}"
  eval "$Command &"

  Multicore_process[Multicore_processcount]=$!
  Multicore_image[Multicore_processcount]="${2:-}"
  Multicore_memory[Multicore_processcount]="${3:-0}"

  return 0
}
multicore_wait() {
  local Process Error=
  for Process in $(seq ${Multicore_maxprocesses:-0}); do
    [ "${Multicore_process[$Process]}" ] && {
      multicore_waitprocess "${Multicore_process[$Process]}" || {
        multicore_break
        Error=1
      }
      [ "$Error" ] && break
      [ "${Multicore_image[$Process]}" ] && showimage "${Multicore_image[$Process]}"
    }
    Multicore_process[$Process]=""
    Multicore_image[$Process]=""
    Multicore_memory[$Process]="0"
  done
  [ "$Error" ] && return 1
  Multicore_processcount=0
  return 0
}
multicore_waitprocess() {
  local Error=
  while sleep 0.2 ; do
    ps -p "${1:-}" >/dev/null || break
    ifcmdbreak && Error=1
    [ "$Error" ] && break
  done
  [ "$Error" ] && return 1
  wait "${1:-}"
  return $?
}
multicore_break() {
  local Process
  for Process in $(seq ${Multicore_maxprocesses:-0}); do
    [ "${Multicore_process[$Process]}" ] && {
      #verbose "multicore_break: Sending SIGINT to $(ps -p ${Multicore_process[$Process]})"
      kill "${Multicore_process[$Process]}"
      wait "${Multicore_process[$Process]}"
      Multicore_process[$Process]=""
      Multicore_image[$Process]=""
      Multicore_memory[$Process]="0"
    }
  done
}
multicore_init() {
  # declare global variables
  local Process
  Multicore_maxprocesses="$(nproc)"
  Multicore_maxprocesses="${Multicore_maxprocesses:-1}"
  Multicore_maxprocesses="$((Multicore_maxprocesses * 2))"
  for Process in $(seq $Multicore_maxprocesses); do
    Multicore_process[$Process]=""
    Multicore_image[$Process]=""
    Multicore_memory[$Process]="0"
  done
  Multicore_processcount=0
  Multicore_minram=250000
  Multicore_maxprocesses=$Multicore_maxprocesses
}
printfreememory() {
  # print current free memory including zram
  local Memory Line Zram
  Memory="$(LC_ALL=C free | grep "Mem:" | LC_ALL=C awk '{print $7}')"
  while read Line; do
    Zram="$(LC_ALL=C awk 'BEGIN {OFMT = "%.0f"} {print ($3 - $4)}' <<< "$Line")"
    Zram="$(LC_ALL=C awk 'BEGIN {OFMT = "%.0f"} {print $1 / 1000}' <<< "$Zram")"
    #Memory="$((Memory + Zram))"
  done < <(/sbin/swapon --bytes | grep zram)
  echo $Memory
}
ifcmdbreak() {
  #[ -e "$Cachedir/exit" ]
  return 1
}

### main
align_focusstack() {
  local Log Line Image Croppedimagelist Options= Command Count
  local X Y W H Lmax=0 Tmax=0 Rmin=10000000 Bmin=10000000
  local Error=

  # align with focus-stack
  [ "$Contrast" = "yes" ]     || Options="$Options --no-contrast"
  [ "$Whitebalance" = "yes" ] || Options="$Options --no-whitebalance"
  echo "stackalign: Aligning $Imagenumber images with focus-stack"
  Log="$(focus-stack --verbose --output="$Destinationdir/" --jpgquality=100 --align-only $Options ${Sourceimagelist[@]})" || return 1

  # calculate smallest valid area common to all aligned images
  while read Line; do
    Line="$(cut -d' ' -f4- <<< "$Line")"
    X="$(digitonly "$(cut -d, -f1 <<< "$Line")" )"
    Y="$(digitonly "$(cut -d, -f2 <<< "$Line")" )"
    W="$(digitonly "$(cut -d, -f3 <<< "$Line")" )"
    H="$(digitonly "$(cut -d, -f4 <<< "$Line")" )"
    [ "$X" -gt "$Lmax" ] && Lmax="$X"
    [ "$Y" -gt "$Tmax" ] && Tmax="$Y"
    [ "$((X+W))" -lt "$Rmin" ] && Rmin="$((X+W))"
    [ "$((Y+H))" -lt "$Bmin" ] && Bmin="$((Y+H))"
  done < <(grep "valid area" <<< "$Log")
  X="$Lmax"
  Y="$Tmax"
  W="$((Rmin-Lmax-1))"
  H="$((Bmin-Tmax-1))"

  # crop aligned images to valid area
  for Count in $(seq $Imagenumber); do
    [ "$Error" ] && break
    Image="$(basename "${Sourceimagelist[$Count]}")"
    Image="$Destinationdir/$Image"
    Croppedimagelist[$Count]="$(rev <<< "$Image" | cut -d. -f2- | rev).${Imageformat}"
    Alignedimagelist[$Count]="$Image"
    printsameline "stackalign: Cropping image $Count / $Imagenumber ${Croppedimagelist[$Count]}"

    Command="
$Magickbin
  ${Alignedimagelist[$Count]}
    -crop ${W}x${H}+${X}+${Y}
    +repage
    -background transparent
    -rotate ${Rotation:-0}
    $Storeimage
    -write ${Croppedimagelist[$Count]}
    -delete 0
  -exit"

    multicore "$(tr -d "\n" <<< "$Command")" "${Croppedimagelist[$Count]}" || Error=1
  done
  echo ""
  multicore_wait || Error=1

  for Count in $(seq $Imagenumber); do
    [ "${Alignedimagelist[$Count]}" != "${Croppedimagelist[$Count]}" ] && rm "${Alignedimagelist[$Count]}"
  done

  return ${Error:-0}
}
declare_variables() {
  # Global variables with default values
  Alignedimagelist=""
  Contrast="no"
  Destinationdir="./aligned"
  Imageformat="tif"
  Imagenumber="0"
  Magickbin=""
  Rotation="0"
  Showimageprocessing=""
  Sourceimagelist=""
  Storeimage="
    -orient undefined
    +repage
    -compress lzw
    -type TrueColorAlpha
    -define png:compression-filter=2
    -define png:compression-level=4
    -define png:compression-strategy=1"
  Whitebalance="no"
  return 0
}
parse_options() {
  local Shortoptions Longoptions Parsedoptions Terminate=

  Shortoptions="hV"
  Longoptions="contrast::,help,output:,rotate:,showimage,whitebalance::"

  Parsedoptions="$(getopt --options "$Shortoptions" --longoptions "$Longoptions" --name "$0" -- "$@")" || error "Error parsing options."
  eval set -- "$Parsedoptions"

  while [ $# -gt 0 ]; do
    case "${1:-}" in
      --contrast)     Contrast="${2:-yes}" ; shift ;;
      -h|--help)      usage ; Terminate="yes" ;;
      --output)       Destinationdir="${2:-$Destinationdir}" ; shift ;;
      --rotate)       Rotation="${2:-0}" ; shift ;;
      -V|--showimage) Showimageprocessing="yes" ;;
      --whitebalance) Whitebalance="${2:-yes}" ; shift ;;
      --) ;;
      *)
        Imagenumber="$((Imagenumber+1))"
        Sourceimagelist[$Imagenumber]="${1:-}"
        [ -f "${1:-}" ] || error "File not found: ${1:-}"
      ;;
    esac
    shift
  done

  [ "$Terminate" ] && exit 0
  return 0
}
finish() {
  trap - SIGINT
  trap - EXIT
  trap - ERR
  multicore_break
  exit ${1:-0}
}
main() {
  set -Eu
  trap finish SIGINT
  trap finish EXIT
  trap finish ERR

  declare_variables
  parse_options "$@" || return 1
  multicore_init

  [ "$Imagenumber" = "0" ] && error "No images specified."

  # check dependencies
  command -v focus-stack >/dev/null || error "focus-stack not found."
  command -v convert >/dev/null && Magickbin="convert"
  command -v magick  >/dev/null && Magickbin="magick"
  [ -z "$Magickbin" ] && error "imagemagick not found."

  # -V, --showimage
  [ "$Showimageprocessing" = "yes" ] && {
    command -v geeqie >/dev/null || error "-V, --showimage: geeqie not found."
    geeqie -t -r --File:"${Sourceimagelist[1]}" 2>/dev/null
  }

  # run
  mkdir -p "$Destinationdir" || return 1
  align_focusstack || return 1

  return 0
}
main "$@"
finish