netron/tools/onnx

#!/bin/bash

set -e

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

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

if [ $(which python3) ] && [ $(which pip3) ]; then
    python="python3"
    pip="pip3"
else
    python="python"
    pip="pip"
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 -p --quiet
        git -C "${third_party}/${1}" reset --quiet --hard origin/master
    else
        echo "Clone ${2}..."
        git -C "${third_party}" clone --recursive ${2}
    fi
    git submodule update --init
}

clean() {
    bold "onnx clean"
    rm -rf ${virtualenv}
    rm -rf ${third_party}/${identifier}
    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"
    virtualenv --quiet -p ${python} ${virtualenv}
    source ${virtualenv}/bin/activate
    export ONNX_ML=1
    export ONNX_NAMESPACE=onnx
    ${pip} install --quiet ${third_party}/onnx
    deactivate
}

schema() {
    bold "onnx schema"
    source ${virtualenv}/bin/activate
    ${node_modules}/protobufjs/bin/pbjs -t static-module -w closure --no-encode --no-delimited --no-comments --keep-case --decode-text -r onnx -o ${src}/onnx-proto.js ${third_party}/${identifier}/onnx/onnx-ml.proto ${third_party}/${identifier}/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
    deactivate
}

metadata() {
    bold "onnx metadata"
    source ${virtualenv}/bin/activate
    pushd ${tools} > /dev/null
    ${python} onnx-script.py metadata
    popd > /dev/null
    deactivate
}

convert() {
    bold "onnx convert"
    source ${virtualenv}/bin/activate
    ${pip} install --quiet sklearn
    ${pip} install --quiet lightgbm
    ${pip} install --quiet ${third_party}/onnxmltools
    ${python} ${tools}/onnx-script.py convert ${1}
    deactivate
}

infer() {
    bold "onnx infer"
    source ${virtualenv}/bin/activate
    ${python} ${tools}/onnx-script.py infer ${1}
    deactivate
}

optimize() {
    bold "onnx optimize"
    source ${virtualenv}/bin/activate
    ${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