#!/bin/bash

#
# Usage:
#
if [[ $# == 0 ]] ; then

    echo
    echo
    echo "Usage:   $(basename $0) <room> <action> option"
    echo
    echo "Example: $(basename $0) boys_room test clean|messy"
    echo "         $(basename $0) boys_room train"
    echo "         $(basename $0) boys_room score [image.jpg]"
    echo
    exit 0

fi

#
# Model selection
#

# InceptionV3
# Most accurate, but also slowest
MODEL_NAME=inceptionv3
INPUT_HEIGHT=""
INPUT_WIDTH=""
TFHUB_MODULE=https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1

# 224px 100% MobileNetV2 = 1.102s
# This model seems accurate and fast
#MODEL_NAME=mobilenetv2_224_100
#INPUT_HEIGHT="--input_height=224"
#INPUT_WIDTH="--input_width=224"
#TFHUB_MODULE=https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/1

# 96px 35% MobileNetV2
# This model wasn't very accurate when I didn't have too much training data
#MODEL_NAME=mobilenetv2_96_035
#INPUT_HEIGHT="--input_height=96"
#INPUT_WIDTH="--input_width=96"
#TFHUB_MODULE=https://tfhub.dev/google/imagenet/mobilenet_v2_035_96/feature_vector/2

#
# Config
#
ROOT=/b/tf
cd $ROOT
MODEL_ROOM=$1
MODEL_ROOT=$ROOT/models/$MODEL_ROOM
DATA_TEST=$ROOT/data_testing/$MODEL_ROOM
DATA_TRAIN=$ROOT/data_training/$MODEL_ROOM
MODEL=$MODEL_ROOT/$MODEL_NAME

#
# Train a new model
#
if [[ $2 == "train" ]] ; then

    python retrain.py --image_dir $DATA_TRAIN --output_graph=$MODEL.pb --output_labels=$MODEL.txt --tfhub_module $TFHUB_MODULE
    exit 0

fi

#
# Test a URL or file
#
if [[ $2 == "score" ]] ; then

    # Determine if they specified a URL or a local image
    #regex='(https?|ftp|file)://[-A-Za-z0-9\+&@#/%?=~_|!:,.;]*[-A-Za-z0-9\+&@#/%=~_|]'
    #if [[ $3 =~ $regex ]]; then

    # If they didn't specify a path-to-url, download from webcam:
    if [[ $# -eq 2 ]]; then

        #
        # Specified a URL, get it to tmp
        #
        if [[ $1 == "boys_room" ]]; then
            URL=http://192.168.1.21/snapshot.cgi
        fi

        if [[ $1 == "kitchen" ]]; then
            URL=http://192.168.1.23/snapshot.cgi
        fi

        if [[ $1 == "upstairs" ]]; then
            URL=http://192.168.1.25/snapshot.cgi
        fi

        TMP_IMAGE=/tmp/how_clean_$MODEL_ROOM.jpg

        # Repeat in case camera failed with zero length / bad image
        wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        FILESIZE=$(stat -c%s "$TMP_IMAGE")
        if [ $FILESIZE -lt 50 ]; then
            sleep 2
            wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        fi
        if [ $FILESIZE -lt 50 ]; then
            sleep 2
            wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        fi
        if [ $FILESIZE -lt 50 ]; then
            sleep 2
            wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        fi
        if [ $FILESIZE -lt 50 ]; then
            sleep 2
            wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        fi
        if [ $FILESIZE -lt 50 ]; then
            sleep 2
            wget "$URL" -O $TMP_IMAGE > /dev/null 2>&1
        fi

    else

        # User specified a path-to-image in arg #3
        TMP_IMAGE=$3

    fi

    #
    # This outputs how "clean" the room is on a scale of 1..100
    #
    { python label_image.py --graph=$MODEL.pb --labels=$MODEL.txt $INPUT_WIDTH $INPUT_HEIGHT --input_layer=Placeholder --output_layer=final_result --image=$TMP_IMAGE; } 2>&1 | grep "clean" | cut -d ' ' -f2 | awk '{printf "%.0f\n", $1*100}'

    exit 0
fi

#
# Loop through 12 images help back for testing:
#
if [[ $2 == "test" ]] ; then
    for i in {1..12}
        do
            { time python label_image.py --graph=$MODEL.pb --labels=$MODEL.txt $INPUT_WIDTH $INPUT_HEIGHT --input_layer=Placeholder --output_layer=final_result --image=$DATA_TEST/$3/$i.jpg; } 2>&1 | grep "$3\|real"
            echo ---
        done
    exit 0
fi