#!/bin/sh -ex

pip install tensorflow-gpu==1.4 --user

train_data=$1
pretrained_model=$2
model_output=$3
gdef_output=$4

if [ "$pretrained_model" != "no_model" ]; then
    mkdir ./pretrained_model
    cd ./pretrained_model
    tar -xzf $pretrained_model
    cd ..
fi

mkdir ./data/
cd ./data/
tar -xzf $train_data
cd ..

export YTF_NIRVANA=1
python train.py

tar -czf $model_output --exclude "*.gdef" ./model-data/
tar -czf $gdef_output ./model-data/*.gdef
