# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="TZSL793i7KuM" papermill={"duration": 0.009793, "end_time": "2021-02-09T19:20:12.335229", "exception": false, "start_time": "2021-02-09T19:20:12.325436", "status": "completed"} tags=[]
# # Setup
#

# + papermill={"duration": 1.697233, "end_time": "2021-02-09T19:20:14.041904", "exception": false, "start_time": "2021-02-09T19:20:12.344671", "status": "completed"} tags=[]
# !pip install python-dotenv

import os

from dotenv import load_dotenv

load_dotenv()

bucket = os.getenv("BUCKET", "")

# + id="hVPzEKoLuEHy" papermill={"duration": 0.017454, "end_time": "2021-02-09T19:20:14.069706", "exception": false, "start_time": "2021-02-09T19:20:14.052252", "status": "completed"} tags=[]
NUM_TRAIN_STEPS = 1
MODEL_TYPE = 'ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18'
CONFIG_TYPE = 'ssd_mobilenet_v1_quantized_300x300_coco14_sync'

import os
CLOUD_ANNOTATIONS_MOUNT = os.path.join('tmp', bucket)
ANNOTATIONS_JSON_PATH   = os.path.join(CLOUD_ANNOTATIONS_MOUNT, '_annotations.json')

CHECKPOINT_PATH = 'tmp/checkpoint'
OUTPUT_PATH     = 'tmp/output'
EXPORTED_PATH   = 'tmp/exported'
DATA_PATH       = 'tmp/data'

LABEL_MAP_PATH    = os.path.join(DATA_PATH, 'label_map.pbtxt')
TRAIN_RECORD_PATH = os.path.join(DATA_PATH, 'train.record')
VAL_RECORD_PATH   = os.path.join(DATA_PATH, 'val.record')


# + [markdown] id="3XINCKkPshgz" papermill={"duration": 0.009653, "end_time": "2021-02-09T19:20:14.088664", "exception": false, "start_time": "2021-02-09T19:20:14.079011", "status": "completed"} tags=[]
# # Install the TensorFlow Object Detection API
# In order to use the TensorFlow Object Detection API, we need to clone it's GitHub Repo.
#
# ### Dependencies
# Most of the dependencies required come preloaded in Google Colab. The only additional package we need to install is TensorFlow.js, which is used for converting our trained model to a model that is compatible for the web.
#
# ### Protocol Buffers
# The TensorFlow Object Detection API relies on what are called `protocol buffers` (also known as `protobufs`). Protobufs are a language neutral way to describe information. That means you can write a protobuf once and then compile it to be used with other languages, like Python, Java or C.
#
# The `protoc` command used below is compiling all the protocol buffers in the `object_detection/protos` folder for Python.
#
# ### Environment
# To use the object detection api we need to add it to our `PYTHONPATH` along with `slim` which contains code for training and evaluating several widely used Convolutional Neural Network (CNN) image classification models.

# + id="o33_jgwGm3NV" papermill={"duration": 12.580909, "end_time": "2021-02-09T19:20:26.679205", "exception": false, "start_time": "2021-02-09T19:20:14.098296", "status": "completed"} tags=[]
import os
import sys
import pathlib

# Clone the tensorflow models repository if it doesn't already exist
if "models" in pathlib.Path.cwd().parts:
  while "models" in pathlib.Path.cwd().parts:
    os.chdir('..')
elif not pathlib.Path('models').exists():
  # !git clone --depth 1 https://github.com/cloud-annotations/models

# !pip install pycocotools==2.0.0 numpy==1.17.5 tensorflow==1.15 tensorflow-hub==0.3.0
# !pip install tf_slim
# !pip install lvis
# !pip install --no-deps tensorflowjs==1.4.0

# %cd models/research
# !protoc object_detection/protos/*.proto --python_out=.

pwd = os.getcwd()
# we need to set both PYTHONPATH for shell scripts and sys.path for python cells
sys.path.append(pwd)
sys.path.append(os.path.join(pwd, 'slim'))
if 'PYTHONPATH' in os.environ:
  os.environ['PYTHONPATH'] += f':{pwd}:{pwd}/slim'
else:
  os.environ['PYTHONPATH'] = f':{pwd}:{pwd}/slim'
# %cd ../..

# + [markdown] id="ISX8k0TfdDHj" papermill={"duration": 0.017993, "end_time": "2021-02-09T19:20:26.715653", "exception": false, "start_time": "2021-02-09T19:20:26.697660", "status": "completed"} tags=[]
# # Generate a Label Map
# One piece of data the Object Detection API needs is a label map protobuf. The label map associates an integer id to the text representation of the label. The ids are indexed by 1, meaning the first label will have an id of 1 not 0.
#
# Here is an example of what a label map looks like:
# ```
# item {
#   id: 1
#   name: 'Cat'
# }
#
# item {
#   id: 2
#   name: 'Dog'
# }
#
# item {
#   id: 3
#   name: 'Gold Fish'
# }
# ```
#

# + id="nJsKCG3UdDsn" papermill={"duration": 0.025642, "end_time": "2021-02-09T19:20:26.758480", "exception": false, "start_time": "2021-02-09T19:20:26.732838", "status": "completed"} tags=[]
import os
import json

# Get a list of labels from the annotations.json
labels = {}
with open(ANNOTATIONS_JSON_PATH) as f:
  annotations = json.load(f)
  labels = annotations['labels']

# Create a file named label_map.pbtxt
os.makedirs(DATA_PATH, exist_ok=True)
with open(LABEL_MAP_PATH, 'w') as f:
  # Loop through all of the labels and write each label to the file with an id
  for idx, label in enumerate(labels):
    f.write('item {\n')
    f.write("\tname: '{}'\n".format(label))
    f.write('\tid: {}\n'.format(idx + 1)) # indexes must start at 1
    f.write('}\n')

# + [markdown] id="siRNKiuvsz25" papermill={"duration": 0.014855, "end_time": "2021-02-09T19:20:26.788320", "exception": false, "start_time": "2021-02-09T19:20:26.773465", "status": "completed"} tags=[]
# # Generate TFRecords
# The TensorFlow Object Detection API expects our data to be in the format of TFRecords.
#
# The TFRecord format is a collection of serialized feature dicts, one for each image, looking something like this:
# ```
# {
#   'image/height': 1800,
#   'image/width': 2400,
#   'image/filename': 'image1.jpg',
#   'image/source_id': 'image1.jpg',
#   'image/encoded': ACTUAL_ENCODED_IMAGE_DATA_AS_BYTES,
#   'image/format': 'jpeg',
#   'image/object/bbox/xmin': [0.7255949630314233, 0.8845598428835489],
#   'image/object/bbox/xmax': [0.9695875693160814, 1.0000000000000000],
#   'image/object/bbox/ymin': [0.5820120073891626, 0.1829972290640394],
#   'image/object/bbox/ymax': [1.0000000000000000, 0.9662484605911330],
#   'image/object/class/text': (['Cat', 'Dog']),
#   'image/object/class/label': ([1, 2])
# }
# ```

# + id="cAkOvP-gZR1x" papermill={"duration": 9.357212, "end_time": "2021-02-09T19:20:36.160724", "exception": false, "start_time": "2021-02-09T19:20:26.803512", "status": "completed"} tags=[]
import os
import io
import json
import random

import PIL.Image
import tensorflow as tf

from object_detection.utils import dataset_util
from object_detection.utils import label_map_util

def create_tf_record(images, annotations, label_map, image_path, output):
  # Create a train.record TFRecord file.
  with tf.python_io.TFRecordWriter(output) as writer:
    # Loop through all the training examples.
    for image_name in images:
      try:
        # Make sure the image is actually a file
        img_path = os.path.join(image_path, image_name)   
        if not os.path.isfile(img_path):
          continue
          
        # Read in the image.
        with tf.gfile.GFile(img_path, 'rb') as fid:
          encoded_jpg = fid.read()

        # Open the image with PIL so we can check that it's a jpeg and get the image
        # dimensions.
        encoded_jpg_io = io.BytesIO(encoded_jpg)
        image = PIL.Image.open(encoded_jpg_io)
        if image.format != 'JPEG':
          raise ValueError('Image format not JPEG')

        width, height = image.size

        # Initialize all the arrays.
        xmins = []
        xmaxs = []
        ymins = []
        ymaxs = []
        classes_text = []
        classes = []

        # The class text is the label name and the class is the id. If there are 3
        # cats in the image and 1 dog, it may look something like this:
        # classes_text = ['Cat', 'Cat', 'Dog', 'Cat']
        # classes      = [  1  ,   1  ,   2  ,   1  ]

        # For each image, loop through all the annotations and append their values.
        for a in annotations[image_name]:
          if ("x" in a and "x2" in a and "y" in a and "y2" in a):
            label = a['label']
            xmins.append(a["x"])
            xmaxs.append(a["x2"])
            ymins.append(a["y"])
            ymaxs.append(a["y2"])
            classes_text.append(label.encode("utf8"))
            classes.append(label_map[label])
       
        # Create the TFExample.
        tf_example = tf.train.Example(features=tf.train.Features(feature={
          'image/height': dataset_util.int64_feature(height),
          'image/width': dataset_util.int64_feature(width),
          'image/filename': dataset_util.bytes_feature(image_name.encode('utf8')),
          'image/source_id': dataset_util.bytes_feature(image_name.encode('utf8')),
          'image/encoded': dataset_util.bytes_feature(encoded_jpg),
          'image/format': dataset_util.bytes_feature('jpeg'.encode('utf8')),
          'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
          'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
          'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
          'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
          'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
          'image/object/class/label': dataset_util.int64_list_feature(classes),
        }))
        if tf_example:
          # Write the TFExample to the TFRecord.
          writer.write(tf_example.SerializeToString())
      except ValueError:
        print('Invalid example, ignoring.')
        pass
      except IOError:
        print("Can't read example, ignoring.")
        pass

with open(ANNOTATIONS_JSON_PATH) as f:
  annotations = json.load(f)['annotations']
  image_files = [image for image in annotations.keys()]
  # Load the label map we created.
  label_map = label_map_util.get_label_map_dict(LABEL_MAP_PATH)

  random.seed(42)
  random.shuffle(image_files)
  num_train = int(0.7 * len(image_files))
  train_examples = image_files[:num_train]
  val_examples = image_files[num_train:]

  create_tf_record(train_examples, annotations, label_map, CLOUD_ANNOTATIONS_MOUNT, TRAIN_RECORD_PATH)
  create_tf_record(val_examples, annotations, label_map, CLOUD_ANNOTATIONS_MOUNT, VAL_RECORD_PATH)

# + [markdown] id="n6DhYpAS7gX2" papermill={"duration": 0.016427, "end_time": "2021-02-09T19:20:36.194706", "exception": false, "start_time": "2021-02-09T19:20:36.178279", "status": "completed"} tags=[]
# # Download a base model
# Training a model from scratch can take days and tons of data. We can mitigate this by using a pretrained model checkpoint. Instead of starting from nothing, we can add to what was already learned with our own data.
#
# There are several pretrained model checkpoints that can be downloaded from the [model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md).
#
# The model we will be training is the SSD MobileNet architecture. SSD MobileNet models have a very small file size and can execute very quickly, compromising little accuracy, which makes it perfect for running in the browser. Additionally, we will be using `quantization`. When we say the model is `quantized` it means instead of using `float32` as the datatype of our numbers we are using `float16` or `int8`.
#
# - `float32(PI)` = `3.1415927` 32 bits
# - `float16(PI)` = `3.14` 16 bits
# - `int8(PI)` = `3` 8 bits
#
# We do this because it can cut our model size down by around a factor of 4! An unquantized version of SSD MobileNet that I trained was `22.3 MB`, but the quantized version was `5.7 MB` that's a `~75%` reduction 🎉

# + id="oHD1Jm0v7jfz" papermill={"duration": 7.179808, "end_time": "2021-02-09T19:20:43.390489", "exception": false, "start_time": "2021-02-09T19:20:36.210681", "status": "completed"} tags=[]
import os
import tarfile

import six.moves.urllib as urllib

download_base = 'http://download.tensorflow.org/models/object_detection/'
model = MODEL_TYPE + '.tar.gz'
tmp = 'tmp/checkpoint.tar.gz'

if not (os.path.exists(CHECKPOINT_PATH)):
  # Download the checkpoint
  opener = urllib.request.URLopener()
  opener.retrieve(download_base + model, tmp)

  # Extract all the `model.ckpt` files.
  with tarfile.open(tmp) as tar:
    for member in tar.getmembers():
      member.name = os.path.basename(member.name)
      if 'model.ckpt' in member.name:
        tar.extract(member, path=CHECKPOINT_PATH)

  os.remove(tmp)

# + [markdown] id="UXlvFvwUHrui" papermill={"duration": 0.015074, "end_time": "2021-02-09T19:20:43.421587", "exception": false, "start_time": "2021-02-09T19:20:43.406513", "status": "completed"} tags=[]
# # Model config
# The final thing we need to do is inject our pipline with the amount of labels we have and where to find the label map, TFRecord and model checkpoint. We also need to change the the batch size, because the default batch size of 128 is too large for Colab to handle.

# + id="C8CVExv6HsJS" papermill={"duration": 0.131021, "end_time": "2021-02-09T19:20:43.567468", "exception": false, "start_time": "2021-02-09T19:20:43.436447", "status": "completed"} tags=[]
import re

from google.protobuf import text_format

from object_detection.utils import config_util
from object_detection.utils import label_map_util

pipeline_skeleton = 'models/research/object_detection/samples/configs/' + CONFIG_TYPE + '.config'
configs = config_util.get_configs_from_pipeline_file(pipeline_skeleton)

label_map = label_map_util.get_label_map_dict(LABEL_MAP_PATH)
num_classes = len(label_map.keys())
meta_arch = configs["model"].WhichOneof("model")

override_dict = {
  'model.{}.num_classes'.format(meta_arch): num_classes,
  'train_config.batch_size': 24,
  'train_input_path': TRAIN_RECORD_PATH,
  'eval_input_path': VAL_RECORD_PATH,
  'train_config.fine_tune_checkpoint': os.path.join(CHECKPOINT_PATH, 'model.ckpt'),
  'label_map_path': LABEL_MAP_PATH
}

configs = config_util.merge_external_params_with_configs(configs, kwargs_dict=override_dict)
pipeline_config = config_util.create_pipeline_proto_from_configs(configs)
config_util.save_pipeline_config(pipeline_config, DATA_PATH)

# + [markdown] id="FNYIZK1xVNAa" papermill={"duration": 0.017126, "end_time": "2021-02-09T19:20:43.601542", "exception": false, "start_time": "2021-02-09T19:20:43.584416", "status": "completed"} tags=[]
# # Start training
# We can start a training run by calling the `model_main` script, passing:
# - The location of the `pipepline.config` we created
# - Where we want to save the model
# - How many steps we want to train the model (the longer you train, the more potential there is to learn)
# - The number of evaluation steps (or how often to test the model) gives us an idea of how well the model is doing

# + id="Wv5h2bwBVO0V" papermill={"duration": 112.735562, "end_time": "2021-02-09T19:22:36.354417", "exception": false, "start_time": "2021-02-09T19:20:43.618855", "status": "completed"} tags=[]
# !rm -rf $OUTPUT_PATH
# !python -m object_detection.model_main \
#     --pipeline_config_path=$DATA_PATH/pipeline.config \
#     --model_dir=$OUTPUT_PATH \
#     --num_train_steps=$NUM_TRAIN_STEPS \
#     --num_eval_steps=100

# + [markdown] id="AwNtvgtdoB-C" papermill={"duration": 0.034056, "end_time": "2021-02-09T19:22:36.422437", "exception": false, "start_time": "2021-02-09T19:22:36.388381", "status": "completed"} tags=[]
# # Export inference graph
# After your model has been trained, you might have a few checkpoints available. A checkpoint is usually emitted every 500 training steps. Each checkpoint is a snapshot of your model at that point in training. In the event that a long running training process crashes, you can pick up at the last checkpoint instead of starting from scratch.
#
# We need to export a checkpoint to a TensorFlow graph proto in order to actually use it. We use regex to find the checkpoint with the highest training step and export it.

# + id="BZgP_FZUoE0d" papermill={"duration": 18.571575, "end_time": "2021-02-09T19:22:55.027258", "exception": false, "start_time": "2021-02-09T19:22:36.455683", "status": "completed"} tags=[]
import os
import re
import json

from object_detection.utils.label_map_util import get_label_map_dict

regex = re.compile(r"model\.ckpt-([0-9]+)\.index")
numbers = [int(regex.search(f).group(1)) for f in os.listdir(OUTPUT_PATH) if regex.search(f)]
TRAINED_CHECKPOINT_PREFIX = os.path.join(OUTPUT_PATH, 'model.ckpt-{}'.format(max(numbers)))

print(f'Using {TRAINED_CHECKPOINT_PREFIX}')

# !rm -rf $EXPORTED_PATH
# !python -m object_detection.export_inference_graph \
#   --pipeline_config_path=$DATA_PATH/pipeline.config \
#   --trained_checkpoint_prefix=$TRAINED_CHECKPOINT_PREFIX \
#   --output_directory=$EXPORTED_PATH

label_map = get_label_map_dict(LABEL_MAP_PATH)
label_array = [k for k in sorted(label_map, key=label_map.get)]

with open(os.path.join(EXPORTED_PATH, 'labels.json'), 'w') as f:
  json.dump(label_array, f)

StarcoderdataJupyterScriptsDedupFiltered