netron/tools/tflite

#!/bin/bash

set -e

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

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

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}" pull --quiet --prune
    else
        git -C "${third_party}" clone --quiet --recursive ${2}
    fi
    git -C "${third_party}/${1}" submodule sync --quiet
    git -C "${third_party}/${1}" submodule update --quiet --init --recursive
}

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 "tflite 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" . > /dev/null
    make > /dev/null
    popd > /dev/null
    [ -n "$(python3 -m pip list --format columns --disable-pip-version-check | grep -w virtualenv)" ] || python3 -m pip install --force-reinstall --user --quiet virtualenv
    [ -d "${virtualenv}" ] || virtualenv --quiet -p python3 ${virtualenv}
}

schema() {
    bold "tflite schema"
    sed 's/namespace tflite;/namespace TFLITE;/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
    mv ./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;
}
EOT
}

visualize() {
    bold "tflite visualize"
    source ${virtualenv}/bin/activate
    python3 -m pip install --quiet tensorflow
    python3 ${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