netron/tools/tflite

#!/bin/bash

set -e

root=$(cd $(dirname ${0})/..; pwd)
build=${root}/build
src=${root}/src
tools=${root}/tools
third_party=${root}/third_party

identifier=tflite
virtualenv=${build}/virtualenv/${identifier}

if [ $(which python3) ]; then
    python="python3"
else
    python="python"
fi

bold() {
    echo "$(tty -s && tput bold)$1$(tty -s && tput sgr0)" 
}

git_sync () {
    mkdir -p "${third_party}"
    if [ -d "${third_party}/${1}" ]; then
        git -C "${third_party}/${1}" fetch --quiet -q
        git -C "${third_party}/${1}" reset --quiet --hard origin/master
    else
        git -C "${third_party}" clone --quiet --recursive ${2}
    fi
}

clean() {
    bold "tflite clean"
    rm -rf ${virtualenv}
    rm -rf ${third_party}/tensorflow
}

sync() {
    bold "tflite sync"
    git_sync flatbuffers https://github.com/google/flatbuffers.git
    git_sync tensorflow https://github.com/tensorflow/tensorflow.git
}

install() {
    bold "flatbuffers install"
    case "$(uname)" in
        "Linux")
            [ -n "$(which cmake)" ] || sudo apt install -y cmake
            ;;
        "Darwin")
            brew list cmake > /dev/null 2>&1 || brew install cmake > /dev/null
            ;;
    esac
    pushd "${third_party}/flatbuffers" > /dev/null
    cmake -G "Unix Makefiles" .
    make
    popd > /dev/null
    [ -n "$(python3 -m pip list --format columns --disable-pip-version-check | grep -w virtualenv)" ] || ${python} -m pip install --force-reinstall --user --quiet virtualenv
    [ -d "${virtualenv}" ] || virtualenv --quiet -p ${python} ${virtualenv}
}

schema() {
    bold "tflite schema"
    sed 's/namespace tflite;/namespace tflite_schema;/g' <${third_party}/tensorflow/tensorflow/lite/schema/schema.fbs >${tools}/tflite.schema.fbs
    ${third_party}/flatbuffers/flatc --no-js-exports --js ${tools}/tflite.schema.fbs
    rm ${tools}/tflite.schema.fbs
    rm ${src}/tflite-schema.js
    cat <<EOT >> ${src}/tflite-schema.js
/* eslint 'indent': [ 'error', 2 ] */
EOT
    cat ./tflite.schema_generated.js >> ${src}/tflite-schema.js
    cat <<EOT >> ${src}/tflite-schema.js
if (typeof module !== 'undefined' && typeof module.exports === 'object') {
  module.exports = tflite_schema;
}
EOT
    rm ./tflite.schema_generated.js
}

visualize() {
    bold "tflite visualize"
    source ${virtualenv}/bin/activate
    ${python} -m pip install --quiet tensorflow
    ${python} ${third_party}/tensorflow/tensorflow/lite/tools/visualize.py $@
    deactivate
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "install") install;;
        "schema") schema;;
        "visualize") visualize ${1} ${2} && shift && shift;;
    esac
done