netrans/netrans_cli/import_model.sh

#!/bin/bash

if [ -z "$NETRANS_PATH" ]; then
    echo "Need to set enviroment variable NETRANS_PATH"
        exit 1
fi

function import_caffe_network()
{
    NAME=$1
    CONVERTCAFFE=$NETRANS_PATH/pnnacc

    CONVERTCAFFE="$CONVERTCAFFE import caffe"

    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi  
    
    echo "=========== Converting $NAME Caffe model ==========="
    if [ -f ${NAME}.caffemodel ]; then
    cmd="$CONVERTCAFFE \
        --model ${NAME}.prototxt \
        --weights ${NAME}.caffemodel \
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data"
    else
    echo "=========== fake Caffe model data file==========="
    cmd="$CONVERTCAFFE \
        --model ${NAME}.prototxt \
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data"
    fi  
}

function import_tensorflow_network()
{
    NAME=$1
    CONVERTF=$NETRANS_PATH/pnnacc

    CONVERTF="$CONVERTF import tensorflow"

    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi  
    
    echo "=========== Converting $NAME Tensorflow model ==========="
    cmd="$CONVERTF \
        --model ${NAME}.pb \
        --output-data ${NAME}.data \
        --output-model ${NAME}.json \
        $(cat inputs_outputs.txt)"
}

function import_onnx_network()
{
    NAME=$1
    CONVERTONNX=$NETRANS_PATH/pnnacc
    CONVERTONNX="$CONVERTONNX import onnx"
    
    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi
    
    echo "=========== Converting $NAME ONNX model ==========="
    cmd="$CONVERTONNX \
        --model  ${NAME}.onnx \
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data"
}

function import_tflite_network()
{
    NAME=$1
    CONVERTTFLITE=$NETRANS_PATH/pnnacc
    CONVERTTFLITE="$CONVERTTFLITE import tflite"
    
    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi  
    
    echo "=========== Converting $NAME TFLite model ==========="
    cmd="$CONVERTTFLITE \
        --model ${NAME}.tflite \
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data"
}

function import_darknet_network()
{
    NAME=$1
    CONVERTDARKNET=$NETRANS_PATH/pnnacc
    CONVERTDARKNET="$CONVERTDARKNET import darknet"
    
    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi
    
    echo "=========== Converting $NAME darknet model ==========="
    cmd="$CONVERTDARKNET \
        --model ${NAME}.cfg \
        --weight ${NAME}.weights \
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data"
}

function import_pytorch_network()
{
    NAME=$1
    CONVERTPYTORCH=$NETRANS_PATH/pnnacc
    CONVERTPYTORCH="$CONVERTPYTORCH import pytorch"
    
    
    if [ -f ${NAME}.json ]; then
        echo -e "\033[31m rm  ${NAME}.json \033[0m" 
        rm ${NAME}.json
    fi
    
    if [ -f ${NAME}.data ]; then
        echo -e "\033[31m rm  ${NAME}.data \033[0m" 
        rm ${NAME}.data
    fi
    
    echo "=========== Converting $NAME pytorch model ==========="
    cmd="$CONVERTPYTORCH \
        --model ${NAME}.pt \q
        --output-model  ${NAME}.json \
        --output-data ${NAME}.data \
	$(cat input_size.txt)"
}

function import_network()
{
    NAME=$(basename "$1")
    pushd $1
    
    if [ -f ${NAME}.prototxt ]; then
        import_caffe_network ${NAME%/}
    elif [ -f ${NAME}.pb ]; then
        import_tensorflow_network ${NAME%/}
    elif [ -f ${NAME}.onnx ]; then
        import_onnx_network ${NAME%/}
    elif [ -f ${NAME}.tflite ]; then
        import_tflite_network ${NAME%/}
    elif [ -f ${NAME}.weights ]; then
        import_darknet_network ${NAME%/}
    elif [ -f ${NAME}.pt ]; then
        import_pytorch_network ${NAME%/}
    else
        echo "=========== can not find suitable model files ==========="
    fi

    echo $cmd
    eval $cmd
    
    if [ -f ${NAME}.data -a -f ${NAME}.json ]; then
        echo -e "\033[31m SUCCESS \033[0m" 
    else
        echo -e "\033[31m ERROR ! \033[0m" 
    fi  
    popd
}

if [ "$#" -ne 1 ]; then
    echo "Input a network name !"
    exit -1
fi

if [ ! -e "${1%/}" ]; then
    echo "Directory ${1%/} does not exist !"
    exit -2
fi

import_network ${1%/}