#!/usr/bin/env bash

# README FIRST:
# ============
# Script to run the TFLite tools benchmark_model or label_image and log results to a file.
#
# Environment prerequisites:
#   1. Export ADB_SERIAL=<your_device_id> in your shell
#   2. Copy files to the device:
#       - binaries to /data/local/tmp/tflite/bin/
#       - models to /data/local/tmp/tflite/models/
#       - test data (sample images, labels file) to /data/local/tmp/tflite/data/
#   3. Setup device CPU/GPU/etc frequencies as appropriate
#
# Saved results location:
#   By default run logs are saved to '[LOGDIR]/[tool_name]-[run_date]/<model_name>.log' where:
#     - [LOGDIR] is set as an env variable, or defaults to working dir
#     - [tool_name] is one of 'benchmark_model' or 'label_image'
#     - [run_date] is the start date-time in the format 'YYYYmmdd_HHMMSS'
#
#   Depending on selected config, the log directory and file names may include one or more of the
#   run options such as - fp16, nnapi, ruy.

set -o pipefail

README_FIRST=yes  # Set this to "yes" disable readme prompt

show_info() {
    echo "
Script to run the TFLite tools benchmark_model or label_image and log results to a file.

Environment prerequisites:
  1. Export ADB_SERIAL=<your_device_id> in your shell
  2. Copy files to the device:
      - binaries to /data/local/tmp/tflite/bin/
      - models to /data/local/tmp/tflite/models/
      - test data (sample images, labels file) to /data/local/tmp/tflite/data/
  3. Setup device CPU/GPU/etc frequencies as appropriate

Saved results location:
  By default run logs are saved to '[LOGDIR]/[tool_name]-[run_date]/<model_name>.log' where:
    - [LOGDIR] is set as an env variable, or defaults to working dir
    - [tool_name] is one of 'benchmark_model' or 'label_image'
    - [run_date] is the start date-time in the format 'YYYYmmdd_HHMMSS'

  Depending on selected config, the log directory and file names may include one or more of the
  run options such as - fp16, nnapi, ruy.
"
}

if [ "$1" = "-h" -o "$1" = "--help" ]; then
    show_info
    exit 0
fi

if [ "$README_FIRST" != "yes" ]; then
    printf "README FIRST:\n============\n"
    show_info
    read -p "Press ENTER key to continue..."
fi

if [ -z "$ADB_SERIAL" ]; then
    printf "\nERROR: Environment ADB_SERIAL must be set!\n"
    exit 1
fi

# Use taskset instead of cpufreq to select processor affinity
if [ "$TASKSET" = "big" ]; then
    taskset_cmd="taskset f0"
elif [ "$TASKSET" = "little" ]; then
    taskset_cmd="taskset 0f"
fi

# Host to which device is connected
ADB_SERVER=${ADB_SERVER:-localhost}
ADB_CMD=(adb -H $ADB_SERVER -s $ADB_SERIAL shell)

warmups=${WARMUPS:-5} # benchmark_model only
iters=${ITERS:-50}
threads=${THREADS:-4}
prof=${PROFILE:-0} # label_image only
nnapi=${NNAPI:-0}
fp16=${FP16:-0}

# Directory for saved logs
[[ "$RUNTOOL" == label* ]] && tool=label_image || tool=benchmark_model
rundate=${RUNDATE:-$(date +%Y%m%d_%H%M%S)}
logpath=${LOGDIR:-$PWD}/$tool-$rundate

# Append config options where necessary
[ $nnapi = 1 ] && logpath=$logpath-nnapi
[ $fp16 = 1 ] && logpath=$logpath-fp16
# [ ! -d "$logpath" ] && mkdir -p $logpath
mkdir -p $(dirname $logpath) && logpath=$logpath.log

# Log runtime config
printf "\nDevice: $(${ADB_CMD[@]} getprop | grep -i ro.product.device)\n" | tee -a $logpath
printf "
Running TFLite with the following params:
  - taskset = ${taskset_cmd:-no}
  - warmups = $warmups runs (benchmark_model)
  - iterations = $iters runs
  - threads = $threads
  - profiling = $prof (label_image)
  - nnapi = $nnapi
  - fp16 = $fp16
  - logpath = $logpath
" | tee -a $logpath

# Folder with tflite installation (tools and data)
BASEDIR=/data/local/tmp/tflite

run_tool() {
    local tool=$1
    local model=$2
    if [ "$tool" = "label_image" ]; then
        # e.g. label_image --labels data/imagenet_labels.txt --image data/grace_hopper.bmp --threads 4 --count 50 --profiling 0 --accelerated 1 --allow_fp16 0 --tflite_model models/float/squeezenet.tflite
        ${ADB_CMD[@]} $taskset_cmd $BASEDIR/bin/label_image -l $BASEDIR/data/imagenet_labels.txt -i $BASEDIR/data/grace_hopper.bmp -t $threads -p $prof -c $iters -a $nnapi -f $fp16 -m $model
    else
        # e.g. benchmark_model --num_threads=4 --warmup_runs=5 --num_runs=50 --use_nnapi=0 --allow_fp16=0 --graph=models/float/squeezenet.tflite
        ${ADB_CMD[@]} $taskset_cmd $BASEDIR/bin/benchmark_model --num_threads=$threads --warmup_runs=$warmups --num_runs=$iters --use_nnapi=$nnapi --allow_fp16=$fp16 --graph=$model
    fi
}

for model in $(${ADB_CMD[@]} find $BASEDIR/models -type f -name "*.tflite"); do
    # logfile=$logpath/$(echo ${model##*/} | sed 's/.tflite//').log # use model name minus path dirname and extension
    run_tool $tool $model 2>&1 | tee -a $logpath
    sleep 1
done

printf "\nRUN COMPLETE: Saved runtime logs to: $logpath\n"