netron/tools/onnx

#!/bin/bash

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

identifier=onnx

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

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

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

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 "onnx clean"
    rm -rf ./third_party/venv/onnx
    rm -rf ./third_party/onnx
    rm -rf ./third_party/onnxmltools
}

sync() {
    bold "onnx sync"
    git_sync onnx https://github.com/onnx/onnx.git
    git_sync onnxmltools https://github.com/onnx/onnxmltools.git
}

install() {
    bold "onnx install"
    case "$(uname)" in
        "Linux")
            [ -n "$(which cmake)" ] || sudo apt install -y cmake
            [ -n "$(which protoc)" ] || sudo apt install -y protobuf-compiler libprotoc-dev
            ;;
        "Darwin")
            brew list cmake > /dev/null 2>&1 || brew install cmake > /dev/null
            brew list protobuf > /dev/null 2>&1 || brew install protobuf > /dev/null
            ;;
    esac
    venv
    ${python} -m pip install --quiet --upgrade protobuf
    export ONNX_ML=1
    export ONNX_NAMESPACE=onnx
    ${python} -m pip install --quiet ./third_party/onnx
    deactivate
}

schema() {
    bold "onnx schema"
    [[ $(grep -U $'\x0D' ./src/onnx-proto.js) ]] && crlf=1
    npx pbjs -t static-module -w closure --no-encode --no-delimited --no-comments --no-convert --no-verify --no-create --keep-case --decode-text -r onnx -o ./src/onnx-proto.js ./third_party/onnx/onnx/onnx-ml.proto ./third_party/onnx/onnx/onnx-operators-ml.proto
    node ./tools/update_pbjs.js array ./src/onnx-proto.js float_data float 1
    node ./tools/update_pbjs.js array ./src/onnx-proto.js double_data double 1
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./src/onnx-proto.js ./src/onnx-proto.js
    fi
}

metadata() {
    bold "onnx metadata"
    [[ $(grep -U $'\x0D' ./src/onnx-metadata.json) ]] && crlf=1
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py metadata
    deactivate
    if [[ -n ${crlf} ]]; then
        unix2dos --quiet --newfile ./src/onnx-metadata.json ./src/onnx-metadata.json
    fi
}

convert() {
    bold "onnx convert"
    venv
    ${python} -m pip install --quiet ./third_party/onnxmltools
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py convert ${1}
    deactivate
}

infer() {
    bold "onnx infer"
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py infer ${1}
    deactivate
}

optimize() {
    bold "onnx optimize"
    venv
    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
    ${python} ./tools/onnx-script.py optimize ${1}
    deactivate
}

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