netron/tools/tf

#!/bin/bash

set -e
pushd $(cd $(dirname ${0})/..; pwd) > /dev/null

case "${OSTYPE}" in
    msys*) python="winpty python";;
    *) python="python";;
esac

venv() {
    env_dir=./third_party/env/tensorflow
    [ -d "${env_dir}" ] || ${python} -m venv ${env_dir}
    case "${OSTYPE}" in
        msys*) source ${env_dir}/Scripts/activate;;
        *) source ${env_dir}/bin/activate;;
    esac
    ${python} -m pip install --quiet --upgrade pip requests
}

clean() {
    echo "tf clean"
    rm -rf "./third_party/env/tensorflow"
    rm -rf "./third_party/source/tensorflow"
    rm -rf "./third_party/source/tflite-support"
}

sync() {
    echo "tf sync"
    [ -d "./third_party/source/tensorflow" ] || git clone --quiet https://github.com/tensorflow/tensorflow.git "./third_party/source/tensorflow"
    git -C "./third_party/source/tensorflow" pull --quiet --prune
    mkdir -p "./third_party/source/tflite-support/tensorflow_lite_support/metadata"
    curl --silent --location --output "./third_party/source/tflite-support/tensorflow_lite_support/metadata/metadata_schema.fbs" "https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/metadata_schema.fbs"
}

install() {
    echo "tf install"
    venv
    ${python} -m pip install --quiet --upgrade wheel
    ${python} -m pip install --quiet --upgrade tf-nightly > /dev/null 2>&1 || true
    deactivate
}

schema() {
    echo "tf schema"
    [[ $(grep -U $'\x0D' ./source/tf-proto.js) ]] && crlf=1
    node ./tools/protoc.js --binary --text --json --root tf --out ./source/tf-proto.js --path ./third_party/source/tensorflow --path ./third_party/source/tensorflow/third_party/xla tensorflow/core/protobuf/saved_model.proto tensorflow/core/protobuf/tensor_bundle.proto tensorflow/core/util/saved_tensor_slice.proto tensorflow/core/util/event.proto tensorflow/core/protobuf/config.proto tensorflow/core/util/memmapped_file_system.proto tensorflow/core/protobuf/fingerprint.proto
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tf-proto.js ./source/tf-proto.js
    fi
    echo "tflite schema"
    [[ $(grep -U $'\x0D' ./source/tflite-schema.js) ]] && crlf=1
    temp1=$(mktemp -d)
    temp2=$(mktemp -d)
    node ./tools/flatc.js --text --root tflite --out ./source/tflite-schema.js ./third_party/source/tensorflow/tensorflow/compiler/mlir/lite/schema/schema.fbs ./third_party/source/tflite-support/tensorflow_lite_support/metadata/metadata_schema.fbs
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tflite-schema.js ./source/tflite-schema.js
    fi
    echo "keras schema"
    [[ $(grep -U $'\x0D' ./source/keras-proto.js) ]] && crlf=1
    node ./tools/protoc.js --binary --text --root tf --out ./source/keras-proto.js --path ./third_party/source/tensorflow --path ./third_party/source/tensorflow/third_party/xla/third_party/tsl tensorflow/python/keras/protobuf/saved_metadata.proto
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/keras-proto.js ./source/keras-proto.js
    fi
}

metadata() {
    echo "tflite metadata"
    export TF_CPP_MIN_LOG_LEVEL=2
    if [[ $(grep -U $'\x0D' ./source/tflite-metadata.json) ]]; then crlf=1; else crlf=; fi
    node ./tools/tflite-script.js
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tflite-metadata.json ./source/tflite-metadata.json
    fi
    echo "tf metadata"
    venv
    if [[ $(grep -U $'\x0D' ./source/tf-metadata.json) ]]; then crlf=1; else crlf=; fi
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    if pip show "tf-nightly" > /dev/null 2>&1; then
        ${python} ./tools/tf_script.py
    fi
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./source/tf-metadata.json ./source/tf-metadata.json
    fi
    deactivate
}

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