# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2020 Google LLC
Licensed under the Apache License, Version 2.0 (the "License")
End of explanation
"""


! pip uninstall tensorflow -y
! pip install tensorflow-gpu==1.15
! pip install -U numpy==1.19.5
! pip install -U pycocotools==2.0.1

import tensorflow as tf
print(tf.__version__)

"""
Explanation: Retrain a detection model for Edge TPU with quant-aware training (TF 1.15)
This notebook uses a set of TensorFlow training scripts to perform transfer-learning on a quantization-aware object detection model and then convert it for compatibility with the Edge TPU.
Specifically, this tutorial shows you how to retrain a MobileNet V1 SSD model so that it detects two pets: Abyssinian cats and American Bulldogs (from the Oxford-IIIT Pets Dataset), using TensorFlow r1.15.
Beware that, compared to a desktop computer, this training can take a lot longer in Colab because Colab provides limited resources for long-running operations. So you'll likely see faster training speeds if you connect this notebook to a local runtime, or instead follow the tutorial to run this training in Docker (which includes more documentation about this process).
Import TensorFlow
End of explanation
"""


! git clone https://github.com/tensorflow/models.git

! cd models && git checkout f788046ca876a8820e05b0b48c1fc2e16b0955bc

! git clone https://github.com/google-coral/tutorials.git

! cp -r tutorials/docker/object_detection/scripts/* models/research/

"""
Explanation: Clone the model and training repos
End of explanation
"""


! apt-get install -y python python-tk
! pip install Cython contextlib2 pillow lxml jupyter matplotlib

# Get protoc 3.0.0, rather than the old version already in the container
! wget https://www.github.com/google/protobuf/releases/download/v3.0.0/protoc-3.0.0-linux-x86_64.zip
! unzip protoc-3.0.0-linux-x86_64.zip -d proto3
! mkdir -p local/bin && mkdir -p local/include
! mv proto3/bin/* local/bin
! mv proto3/include/* local/include
! rm -rf proto3 protoc-3.0.0-linux-x86_64.zip

# Install pycocoapi
! git clone --depth 1 https://github.com/cocodataset/cocoapi.git
! (cd cocoapi/PythonAPI && make -j8)
! cp -r cocoapi/PythonAPI/pycocotools/ models/research/
! rm -rf cocoapi

# Run protoc on the object detection repo (generate .py files from .proto)
% cd models/research/
! ../../local/bin/protoc object_detection/protos/*.proto --python_out=.

import os
os.environ['PYTHONPATH'] += ":/content/models/research:/content/models/research/slim"

"""
Explanation: Import dependencies
For details, see https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md
End of explanation
"""


! python object_detection/builders/model_builder_test.py

"""
Explanation: Just to verify everything is correctly set up:
End of explanation
"""


! ./prepare_checkpoint_and_dataset.sh --network_type mobilenet_v1_ssd --train_whole_model false

"""
Explanation: Convert training data to TFRecord
To train with different images, read how to configure your own training data.
End of explanation
"""


%env NUM_TRAINING_STEPS=500
%env NUM_EVAL_STEPS=100

# If you're retraining the whole model, we suggest thes values:
# %env NUM_TRAINING_STEPS=50000
# %env NUM_EVAL_STEPS=2000

! ./retrain_detection_model.sh --num_training_steps $NUM_TRAINING_STEPS --num_eval_steps $NUM_EVAL_STEPS

"""
Explanation: Perform transfer-learning
The following script takes several hours to finish in Colab. (You can shorten by reducing the steps, but that reduces the final accuracy.)
If you didn't already select "Run all" then you should run all remaining cells now. That will ensure the rest of the notebook completes while you are away, avoiding the chance that the Colab runtime times-out and you lose the training data before you download the model.
End of explanation
"""


! ./convert_checkpoint_to_edgetpu_tflite.sh --checkpoint_num $NUM_TRAINING_STEPS

! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -

! echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list

! sudo apt-get update

! sudo apt-get install edgetpu-compiler	

%cd learn_pet/models/

! ls

! edgetpu_compiler output_tflite_graph.tflite

"""
Explanation: As training progresses, you can see new checkpoint files appear in the models/research/learn_pet/train/ directory.
Compile for the Edge TPU
End of explanation
"""


from google.colab import files

files.download('output_tflite_graph_edgetpu.tflite')
files.download('labels.txt')

"""
Explanation: Download the files:
End of explanation
"""


