Namburger

# MAKE THIS DOCKERFILE

FROM ubuntu:20.04
ENV TZ=US/Central
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

WORKDIR /home
ENV HOME /home
VOLUME /data
EXPOSE 8888

Namburger

# Import tensorflow 1.x and install tf_slim.
%tensorflow_version 1.x
!pip install tf_slim
!pip show tensorflow

# Install protobuf-compiler and the tensorflow's Object Detection API.
!apt-get install protobuf-compiler
!git clone https://github.com/tensorflow/models.git

%cd models/research

Namburger

# Now let's download our training dataset.
%mkdir /content/dataset
%cd /content/dataset
!wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
!wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz
!tar zxf images.tar.gz
!tar zxf annotations.tar.gz

# Only picking Abyssinian Cat and the American Bulldog
# If you wish to train the model on all classes, simply skip this entire cell.

Namburger

# Now we can create the tfrecord files.
%cd /content/models/research
!cp object_detection/data/pet_label_map.pbtxt /content/dataset
!python3 object_detection/dataset_tools/create_pet_tf_record.py \
    --label_map_path="/content/dataset/pet_label_map.pbtxt" \
    --data_dir="/content/dataset" \
    --output_dir="/content/dataset"

Namburger

# Now let's download our ssdlite mobiledet pretrained model from tensorflow's model zoo.
!mkdir /content/pretrained_model
%cd /content/pretrained_model
!wget http://download.tensorflow.org/models/object_detection/ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19.tar.gz
!tar xvf ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19.tar.gz

Namburger

# Edit Pipeline config to load in our new tfrecord that we just created and add quantization aware training.
import tensorflow as tf
from google.protobuf import text_format
from object_detection.protos import pipeline_pb2

# Hack to find out if you have colab pro or not :)
gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
print(gpu_info)
gpu_name = !nvidia-smi --query-gpu=gpu_name --format=csv

Namburger

model {
  ssd {
    num_classes: 2
    image_resizer {
      fixed_shape_resizer {
        height: 320
        width: 320
      }
    }
    feature_extractor {

Namburger

%cd /content
!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
!unzip -o ngrok-stable-linux-amd64.zip

# Starts tensorboard, so we can monitor the training process.
get_ipython().system_raw(
    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'
    .format('/content/train')
)
get_ipython().system_raw('./ngrok http 6006 &')

Namburger

# Make inference graph.
!python3 /content/models/research/object_detection/export_inference_graph.py \
    --input_type=image_tensor \
    --pipeline_config_path=/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config \
    --output_directory=/content/inference_graph \
    --trained_checkpoint_prefix=/content/train/model.ckpt-25000 # Make sure to change this checkpoint to the corresponding num step you set from above.

Namburger

# Let's begin training, expects to take a few hours, time for a good stretch :)
%cd /content/models/research/
!python3 object_detection/model_main.py \
    --logtostderr=true \
    --model_dir=/content/train \
    --pipeline_config_path=/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config