├── .dockerignore ├── .gitignore ├── README.md ├── docker ├── Dockerfile.training └── Dockerfile.training.gpu ├── python ├── calculate_map_graph.py ├── create_data_tf_record.py ├── create_label_map.py └── update_config.py └── run.sh /.dockerignore: -------------------------------------------------------------------------------- 1 | data -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | data 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Object Detection Using Tensorflow on the Raspberry Pi 2 | 3 | Script for object detection from training new model on dataset to exporting quantized graph 4 | 5 | ## Step 1. Setup 6 | 7 | ### Using docker registry 8 | This is the fastest way to use the repo 9 | ``` 10 | # For cpu 11 | docker pull docker.nanonets.com/pi_training 12 | # For gpu 13 | docker pull docker.nanonets.com/pi_training:gpu 14 | ``` 15 | OR 16 | 17 | ### Building locally 18 | #### Docker build script 19 | Should run this script from repository root 20 | ``` 21 | # For cpu 22 | docker build -t pi_training -f docker/Dockerfile.training . 23 | docker image tag pi_training docker.nanonets.com/pi_training 24 | 25 | # For gpu 26 | docker build -t pi_training:gpu -f docker/Dockerfile.training.gpu . 27 | docker image tag pi_training:gpu docker.nanonets.com/pi_training:gpu 28 | ``` 29 | ------ 30 | 31 | ## Step 2. Preparing dataset 32 | Dataset for object detection consists of images of objects you want to detect and annotations which are xml files with coordinates of objects inside images in Pascal VOC format. If you have collected images, you can use tool like [LabelImg](https://github.com/tzutalin/labelImg) to create dataset. 33 | 34 | Copy dataset with `images` folder containing all training images and `annotations` folder containing all respective annotations inside `data` folder in repo which will be mounted by docker as volume 35 | 36 | ## Step 3. Starting training 37 | Tensorboard will be started at port 8000 and run in background 38 | You can specify -h parameter to get help for docker script 39 | 40 | If you have a GPU instance, you need to install [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) 41 | 42 | ``` 43 | # For cpu 44 | sudo docker run -p 8000:8000 -v `pwd`/data:/data docker.nanonets.com/pi_training -m train -a ssd_mobilenet_v1_coco -e ssd_mobilenet_v1_coco_0 -p '{"batch_size":8,"learning_rate":0.003}' 45 | # For gpu 46 | sudo nvidia-docker run -p 8000:8000 -v `pwd`/data:/data docker.nanonets.com/pi_training:gpu -m train -a ssd_mobilenet_v1_coco -e ssd_mobilenet_v1_coco_0 -p '{"batch_size":8,"learning_rate":0.003}' 47 | ``` 48 | 49 | ### Usage 50 | The docker instance on startup runs a script run.sh which takes the following parameters: 51 | ``` 52 | run.sh [-m mode] [-a architecture] [-h help] [-e experiment_id] [-c checkpoint] [-p hyperparameters] 53 | ``` 54 | -h display this help and exit 55 | -m mode: should be either `train` or `export` 56 | -p key value pairs of hyperparameters as json string 57 | -e experiment id. Used as path inside data folder to run current experiment 58 | -c applicable when mode is export, used to specify checkpoint to use for export 59 | 60 | **List of Models (that can be passed to -a):** 61 | 1. ssd_mobilenet_v1_coco 62 | 2. ssd_inception_v2_coco 63 | 3. faster_rcnn_inception_v2_coco 64 | 4. faster_rcnn_resnet50_coco 65 | 5. rfcn_resnet101_coco 66 | 6. faster_rcnn_resnet101_coco 67 | 7. faster_rcnn_inception_resnet_v2_atrous_coco 68 | 8. faster_rcnn_nas 69 | 70 | **Possible hyperparameters to override from -p command in json** 71 | 72 | | Name | Type | 73 | |-----------|-----------------| 74 | | learning_rate | float | 75 | | batch_size | int | 76 | | train_steps | int | 77 | | eval_steps | int | 78 | 79 | ------ 80 | 81 | ## Step 4. Exporting trained model 82 | This command would export trained model in quantized graph that can be used for prediction. You need to specify one of the trained checkpoints from experiment directory that you want to use for prediction with -c command as follows: 83 | 84 | ``` 85 | # For cpu 86 | sudo docker run -v `pwd`/data:/data docker.nanonets.com/pi_training -m export -a ssd_mobilenet_v1_coco -e ssd_mobilenet_v1_coco_0 -c /data/0/model.ckpt-8998 87 | 88 | # For gpu 89 | sudo nvidia-docker run -v `pwd`/data:/data docker.nanonets.com/pi_training:gpu -m export -a ssd_mobilenet_v1_coco -e ssd_mobilenet_v1_coco_0 -c /data/0/model.ckpt-8998 90 | ``` 91 | 92 | Once your done training the model and have exported it you can move this onto a client device like the Raspberry Pi. 93 | For details of how to use on the Raspberry Pi click see https://github.com/NanoNets/TF-OD-Pi-Test 94 | -------------------------------------------------------------------------------- /docker/Dockerfile.training: -------------------------------------------------------------------------------- 1 | FROM tensorflow/tensorflow:1.6.0 2 | 3 | # Setup environment for tensorflow models 4 | RUN apt-get update && apt-get install -y \ 5 | build-essential \ 6 | curl \ 7 | git \ 8 | libfreetype6-dev \ 9 | libpng12-dev \ 10 | libzmq3-dev \ 11 | pkg-config \ 12 | python-dev \ 13 | python-numpy \ 14 | python-pip \ 15 | protobuf-compiler \ 16 | python-pil \ 17 | python-lxml \ 18 | python-tk \ 19 | software-properties-common \ 20 | swig \ 21 | zip \ 22 | zlib1g-dev \ 23 | libcurl3-dev \ 24 | && \ 25 | apt-get clean && \ 26 | rm -rf /var/lib/apt/lists/* 27 | 28 | RUN apt-get install python-pip python-dev build-essential 29 | 30 | # Set up grpc 31 | RUN pip install enum34 futures mock six matplotlib jupyter && \ 32 | pip install --pre 'protobuf>=3.0.0a3' && \ 33 | pip install -i https://testpypi.python.org/simple --pre grpcio 34 | 35 | WORKDIR / 36 | RUN git clone -b r1.4 https://github.com/tensorflow/tensorflow.git 37 | 38 | # Set up Bazel. 39 | 40 | # Running bazel inside a `docker build` command causes trouble, cf: 41 | # https://github.com/bazelbuild/bazel/issues/134 42 | # The easiest solution is to set up a bazelrc file forcing --batch. 43 | RUN echo "startup --batch" >>/etc/bazel.bazelrc 44 | # Similarly, we need to workaround sandboxing issues: 45 | # https://github.com/bazelbuild/bazel/issues/418 46 | RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \ 47 | >>/etc/bazel.bazelrc 48 | # Install the most recent bazel release. 49 | ENV BAZEL_VERSION 0.5.4 50 | WORKDIR / 51 | RUN mkdir /bazel && \ 52 | cd /bazel && \ 53 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 54 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \ 55 | chmod +x bazel-*.sh && \ 56 | ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 57 | cd / && \ 58 | rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh 59 | 60 | WORKDIR /tensorflow 61 | RUN tensorflow/tools/ci_build/builds/configured CPU \ 62 | bazel build tensorflow/tools/graph_transforms:transform_graph 63 | 64 | WORKDIR / 65 | RUN git clone https://github.com/tensorflow/models.git 66 | RUN git clone https://github.com/cocodataset/cocoapi.git 67 | WORKDIR /cocoapi/PythonAPI 68 | RUN pip install Cython && make 69 | RUN cp -r pycocotools /models/research/ 70 | 71 | WORKDIR /models/research/ 72 | RUN protoc object_detection/protos/*.proto --python_out=. 73 | ENV PYTHONPATH=$PYTHONPATH:/models/research/:/models/research/slim 74 | 75 | # Add dataset into docker file 76 | # Add volumes for output graph and 77 | VOLUME /data 78 | 79 | # Add files to dockerfile 80 | ADD . / 81 | RUN chmod +x /run.sh 82 | 83 | EXPOSE 8000 84 | ENTRYPOINT ["/run.sh"] 85 | CMD [""] -------------------------------------------------------------------------------- /docker/Dockerfile.training.gpu: -------------------------------------------------------------------------------- 1 | FROM tensorflow/tensorflow:1.6.0-gpu 2 | 3 | # Setup environment for tensorflow models 4 | RUN apt-get update && apt-get install -y \ 5 | build-essential \ 6 | curl \ 7 | git \ 8 | libfreetype6-dev \ 9 | libpng12-dev \ 10 | libzmq3-dev \ 11 | pkg-config \ 12 | python-dev \ 13 | python-numpy \ 14 | python-pip \ 15 | protobuf-compiler \ 16 | python-pil \ 17 | python-lxml \ 18 | python-tk \ 19 | software-properties-common \ 20 | swig \ 21 | zip \ 22 | zlib1g-dev \ 23 | libcurl3-dev \ 24 | && \ 25 | apt-get clean && \ 26 | rm -rf /var/lib/apt/lists/* 27 | 28 | RUN apt-get install python-pip python-dev build-essential 29 | 30 | #RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \ 31 | # python get-pip.py && \ 32 | # rm get-pip.py 33 | 34 | # Set up grpc 35 | RUN pip install enum34 futures mock six matplotlib jupyter && \ 36 | pip install --pre 'protobuf>=3.0.0a3' && \ 37 | pip install -i https://testpypi.python.org/simple --pre grpcio 38 | 39 | WORKDIR / 40 | RUN git clone -b r1.4 https://github.com/tensorflow/tensorflow.git 41 | 42 | # Set up Bazel. 43 | 44 | # Running bazel inside a `docker build` command causes trouble, cf: 45 | # https://github.com/bazelbuild/bazel/issues/134 46 | # The easiest solution is to set up a bazelrc file forcing --batch. 47 | RUN echo "startup --batch" >>/etc/bazel.bazelrc 48 | # Similarly, we need to workaround sandboxing issues: 49 | # https://github.com/bazelbuild/bazel/issues/418 50 | RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \ 51 | >>/etc/bazel.bazelrc 52 | # Install the most recent bazel release. 53 | ENV BAZEL_VERSION 0.5.4 54 | WORKDIR / 55 | RUN mkdir /bazel && \ 56 | cd /bazel && \ 57 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 58 | curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \ 59 | chmod +x bazel-*.sh && \ 60 | ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \ 61 | cd / && \ 62 | rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh 63 | 64 | WORKDIR /tensorflow 65 | RUN tensorflow/tools/ci_build/builds/configured CPU \ 66 | bazel build tensorflow/tools/graph_transforms:transform_graph 67 | 68 | WORKDIR / 69 | RUN git clone https://github.com/tensorflow/models.git 70 | RUN git clone https://github.com/cocodataset/cocoapi.git 71 | WORKDIR /cocoapi/PythonAPI 72 | RUN pip install Cython && make 73 | RUN cp -r pycocotools /models/research/ 74 | 75 | WORKDIR /models/research/ 76 | RUN protoc object_detection/protos/*.proto --python_out=. 77 | ENV PYTHONPATH=$PYTHONPATH:/models/research/:/models/research/slim 78 | 79 | # Add dataset into docker file 80 | # Add volumes for output graph and 81 | VOLUME /data 82 | 83 | # Add files to dockerfile 84 | ADD . / 85 | RUN chmod +x /run.sh 86 | 87 | EXPOSE 8000 88 | ENTRYPOINT ["/run.sh"] 89 | CMD [""] -------------------------------------------------------------------------------- /python/calculate_map_graph.py: -------------------------------------------------------------------------------- 1 | """ 2 | Script to calculate MAP with quantzed or frozen graph 3 | 4 | """ 5 | import logging 6 | 7 | from object_detection import eval_util 8 | from object_detection.utils import object_detection_evaluation 9 | 10 | # Create evaluation for all test images 11 | 12 | # Tensorflow op so cant use with our graph: convert to result dict with eval_util.result_dict_for_single_example 13 | evaluator.add_single_ground_truth_image_info( 14 | image_id=batch, groundtruth_dict=result_dict) 15 | evaluator.add_single_detected_image_info( 16 | image_id=batch, detections_dict=result_dict) 17 | metrics = evaluator.evaluate() 18 | evaluator.clear() -------------------------------------------------------------------------------- /python/create_data_tf_record.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | r"""Convert the Oxford pet dataset to TFRecord for object_detection. 17 | 18 | See: O. M. Parkhi, A. Vedaldi, A. Zisserman, C. V. Jawahar 19 | Cats and Dogs 20 | IEEE Conference on Computer Vision and Pattern Recognition, 2012 21 | http://www.robots.ox.ac.uk/~vgg/data/pets/ 22 | 23 | Example usage: 24 | ./create_pet_tf_record --data_dir=/home/user/pet \ 25 | --output_dir=/home/user/pet/output 26 | """ 27 | 28 | import hashlib 29 | import io 30 | import logging 31 | import os 32 | import random 33 | import re 34 | 35 | from lxml import etree 36 | import PIL.Image 37 | import tensorflow as tf 38 | 39 | from object_detection.utils import dataset_util 40 | from object_detection.utils import label_map_util 41 | 42 | flags = tf.app.flags 43 | flags.DEFINE_string('data_dir', '', 'Root directory to raw pet dataset.') 44 | flags.DEFINE_string('output_dir', '', 'Path to directory to output TFRecords.') 45 | flags.DEFINE_string('label_map_path', 'data/pet_label_map.pbtxt', 46 | 'Path to label map proto') 47 | FLAGS = flags.FLAGS 48 | 49 | 50 | def get_class_name_from_filename(file_name): 51 | """Gets the class name from a file. 52 | 53 | Args: 54 | file_name: The file name to get the class name from. 55 | ie. "american_pit_bull_terrier_105.jpg" 56 | 57 | Returns: 58 | A string of the class name. 59 | """ 60 | match = re.match(r'([A-Za-z_]+)(_[0-9]+\.jpg)', file_name, re.I) 61 | return match.groups()[0] 62 | 63 | 64 | def dict_to_tf_example(data, 65 | label_map_dict, 66 | image_subdirectory, 67 | ignore_difficult_instances=False): 68 | """Convert XML derived dict to tf.Example proto. 69 | 70 | Notice that this function normalizes the bounding box coordinates provided 71 | by the raw data. 72 | 73 | Args: 74 | data: dict holding PASCAL XML fields for a single image (obtained by 75 | running dataset_util.recursive_parse_xml_to_dict) 76 | label_map_dict: A map from string label names to integers ids. 77 | image_subdirectory: String specifying subdirectory within the 78 | Pascal dataset directory holding the actual image data. 79 | ignore_difficult_instances: Whether to skip difficult instances in the 80 | dataset (default: False). 81 | 82 | Returns: 83 | example: The converted tf.Example. 84 | 85 | Raises: 86 | ValueError: if the image pointed to by data['filename'] is not a valid JPEG 87 | """ 88 | img_path = os.path.join(image_subdirectory, os.path.basename(data['filename'])) 89 | with tf.gfile.GFile(img_path, 'rb') as fid: 90 | encoded_jpg = fid.read() 91 | encoded_jpg_io = io.BytesIO(encoded_jpg) 92 | image = PIL.Image.open(encoded_jpg_io) 93 | 94 | if image.format != 'JPEG': 95 | logging.warning("Image format not jpeg, trying to convert to jpeg") 96 | try: 97 | image = image.convert('RGB') 98 | except: 99 | logging.exception("error converting to jpeg") 100 | return None 101 | 102 | key = hashlib.sha256(encoded_jpg).hexdigest() 103 | 104 | width = int(data['size']['width']) 105 | height = int(data['size']['height']) 106 | 107 | if width < 0 or height < 0: 108 | width, height = image.size 109 | 110 | xmin = [] 111 | ymin = [] 112 | xmax = [] 113 | ymax = [] 114 | classes = [] 115 | classes_text = [] 116 | truncated = [] 117 | poses = [] 118 | difficult_obj = [] 119 | 120 | if not 'object' in data: 121 | return 122 | 123 | for obj in data['object']: 124 | xmin.append(float(obj['bndbox']['xmin']) / width) 125 | ymin.append(float(obj['bndbox']['ymin']) / height) 126 | xmax.append(float(obj['bndbox']['xmax']) / width) 127 | ymax.append(float(obj['bndbox']['ymax']) / height) 128 | classes_text.append(obj['name'].encode('utf8')) 129 | classes.append(label_map_dict[obj['name']]) 130 | 131 | example = tf.train.Example(features=tf.train.Features(feature={ 132 | 'image/height': dataset_util.int64_feature(height), 133 | 'image/width': dataset_util.int64_feature(width), 134 | 'image/filename': dataset_util.bytes_feature( 135 | data['filename'].encode('utf8')), 136 | 'image/source_id': dataset_util.bytes_feature( 137 | data['filename'].encode('utf8')), 138 | 'image/key/sha256': dataset_util.bytes_feature(key.encode('utf8')), 139 | 'image/encoded': dataset_util.bytes_feature(encoded_jpg), 140 | 'image/format': dataset_util.bytes_feature('jpeg'.encode('utf8')), 141 | 'image/object/bbox/xmin': dataset_util.float_list_feature(xmin), 142 | 'image/object/bbox/xmax': dataset_util.float_list_feature(xmax), 143 | 'image/object/bbox/ymin': dataset_util.float_list_feature(ymin), 144 | 'image/object/bbox/ymax': dataset_util.float_list_feature(ymax), 145 | 'image/object/class/text': dataset_util.bytes_list_feature(classes_text), 146 | 'image/object/class/label': dataset_util.int64_list_feature(classes), 147 | 'image/object/difficult': dataset_util.int64_list_feature(difficult_obj), 148 | 'image/object/truncated': dataset_util.int64_list_feature(truncated), 149 | 'image/object/view': dataset_util.bytes_list_feature(poses), 150 | })) 151 | return example 152 | 153 | 154 | def create_tf_record(output_filename, 155 | label_map_dict, 156 | annotations_dir, 157 | image_dir, 158 | examples): 159 | """Creates a TFRecord file from examples. 160 | 161 | Args: 162 | output_filename: Path to where output file is saved. 163 | label_map_dict: The label map dictionary. 164 | annotations_dir: Directory where annotation files are stored. 165 | image_dir: Directory where image files are stored. 166 | examples: Examples to parse and save to tf record. 167 | """ 168 | writer = tf.python_io.TFRecordWriter(output_filename) 169 | for idx, example in enumerate(examples): 170 | if idx % 100 == 0: 171 | logging.info('On image %d of %d', idx, len(examples)) 172 | path = os.path.join(annotations_dir, example) 173 | 174 | if not os.path.exists(path): 175 | logging.warning('Could not find %s, ignoring example.', path) 176 | continue 177 | try: 178 | with tf.gfile.GFile(path, 'r') as fid: 179 | xml_str = fid.read() 180 | xml = etree.fromstring(xml_str) 181 | data = dataset_util.recursive_parse_xml_to_dict(xml)['annotation'] 182 | 183 | tf_example = dict_to_tf_example(data, label_map_dict, image_dir) 184 | if tf_example: 185 | writer.write(tf_example.SerializeToString()) 186 | except Exception as e: 187 | logging.exception("Could not convert example in rf record") 188 | 189 | writer.close() 190 | 191 | 192 | # TODO: Add test for pet/PASCAL main files. 193 | def main(_): 194 | data_dir = FLAGS.data_dir 195 | label_map_dict = label_map_util.get_label_map_dict(FLAGS.label_map_path) 196 | 197 | logging.info('Reading dataset.') 198 | image_dir = os.path.join(data_dir, 'images') 199 | annotations_dir = os.path.join(data_dir, 'annotations') 200 | examples_list = [] 201 | for example in os.listdir(annotations_dir): 202 | examples_list.append(example) 203 | 204 | # Test images are not included in the downloaded data set, so we shall perform 205 | # our own split. 206 | random.seed(42) 207 | random.shuffle(examples_list) 208 | num_examples = len(examples_list) 209 | num_train = int(0.8 * num_examples) 210 | train_examples = examples_list[:num_train] 211 | val_examples = examples_list[num_train:] 212 | logging.info('%d training and %d validation examples.', 213 | len(train_examples), len(val_examples)) 214 | 215 | train_output_path = os.path.join(FLAGS.output_dir, 'train.record') 216 | val_output_path = os.path.join(FLAGS.output_dir, 'val.record') 217 | create_tf_record(train_output_path, label_map_dict, annotations_dir, 218 | image_dir, train_examples) 219 | create_tf_record(val_output_path, label_map_dict, annotations_dir, 220 | image_dir, val_examples) 221 | 222 | if __name__ == '__main__': 223 | tf.app.run() 224 | -------------------------------------------------------------------------------- /python/create_label_map.py: -------------------------------------------------------------------------------- 1 | r"""Create label map from dataset 2 | 3 | 4 | Example usage: 5 | ./create_label_map.py --data_dir=/home/data/ \ 6 | --label_map_path=/home/data/label_map.pbtxt 7 | """ 8 | 9 | import os 10 | import io 11 | import logging 12 | from collections import defaultdict 13 | 14 | from lxml import etree 15 | import PIL.Image 16 | 17 | from google.protobuf import text_format 18 | from object_detection.protos import string_int_label_map_pb2 19 | from object_detection.protos import string_int_label_map_pb2 20 | import tensorflow as tf 21 | 22 | from object_detection.utils import dataset_util 23 | 24 | flags = tf.app.flags 25 | flags.DEFINE_string('data_dir', '', 'Root directory to raw pet dataset.') 26 | flags.DEFINE_string('label_map_path', 'data/label_map.pbtxt', 27 | 'Path to label map proto') 28 | FLAGS = flags.FLAGS 29 | 30 | def get_class_set(annotations_dir, image_dir): 31 | categories = defaultdict(int) 32 | for filename in os.listdir(annotations_dir): 33 | try: 34 | with tf.gfile.GFile(os.path.join(annotations_dir, filename), 'r') as fid: 35 | xml_str = fid.read() 36 | logging.info("xml: ", xml_str) 37 | xml = etree.fromstring(xml_str) 38 | data = dataset_util.recursive_parse_xml_to_dict(xml)['annotation'] 39 | img_path = os.path.join(image_dir, os.path.basename(data['filename'])) 40 | logging.info("image path: ", img_path) 41 | with tf.gfile.GFile(img_path, 'rb') as fid: 42 | encoded_jpg = fid.read() 43 | encoded_jpg_io = io.BytesIO(encoded_jpg) 44 | image = PIL.Image.open(encoded_jpg_io) 45 | if image.format != 'JPEG': 46 | raise ValueError('Image format not JPEG') 47 | for obj in data['object']: 48 | categories[obj['name']] += 1 49 | except Exception as e: 50 | logging.exception('Could not decode xml') 51 | return categories 52 | 53 | 54 | def write_label_map(categories, label_map_path): 55 | label_map = string_int_label_map_pb2.StringIntLabelMap() 56 | label_map_items = [] 57 | for i, category in enumerate(categories): 58 | idx = i + 1 59 | proto = string_int_label_map_pb2.StringIntLabelMapItem() 60 | proto.id = idx 61 | proto.name = category 62 | 63 | label_map_items.append(proto) 64 | 65 | label_map.item.extend(label_map_items) 66 | label_map_str = text_format.MessageToString(label_map) 67 | with open(label_map_path, 'w') as label_map_file: 68 | label_map_file.write(label_map_str) 69 | 70 | 71 | def main(_): 72 | data_dir = FLAGS.data_dir 73 | 74 | logging.info('Creating label map from dataset') 75 | annotations_dir = os.path.join(data_dir, 'annotations') 76 | image_dir = os.path.join(data_dir, 'images') 77 | 78 | categories = get_class_set(annotations_dir, image_dir) 79 | logging.info("No of objects per categories: ", categories) 80 | 81 | logging.info("Creating label map proto file") 82 | write_label_map(categories, FLAGS.label_map_path) 83 | 84 | if __name__ == '__main__': 85 | tf.app.run() -------------------------------------------------------------------------------- /python/update_config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import logging 4 | import tarfile 5 | from six.moves import urllib 6 | import json 7 | import tensorflow as tf 8 | from google.protobuf import text_format 9 | from object_detection.protos import string_int_label_map_pb2 10 | from object_detection.utils import config_util 11 | 12 | flags = tf.app.flags 13 | flags.DEFINE_string('architecture', '', 'Name of architecture') 14 | flags.DEFINE_string('experiment_id', '', 'Id of current experiment output') 15 | flags.DEFINE_string('label_map_path', 'data/pet_label_map.pbtxt', 16 | 'Path to label map proto') 17 | flags.DEFINE_string('data_dir', 'data/', 18 | 'Path to label map proto') 19 | flags.DEFINE_string('hparams', '', 20 | 'Params in json') 21 | FLAGS = flags.FLAGS 22 | 23 | # Map of architecture to configs and urls 24 | arch_map = { 25 | 'ssd_mobilenet_v1_coco': { 26 | 'config': 'ssd_mobilenet_v1_coco.config', 27 | 'url': 'http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2017_11_17.tar.gz', 28 | 'checkpoint': 'ssd_mobilenet_v1_coco_2017_11_17' 29 | }, 30 | 'ssd_inception_v2_coco': { 31 | 'config': 'ssd_inception_v2_coco.config', 32 | 'url': 'http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2017_11_17.tar.gz', 33 | 'checkpoint': 'ssd_inception_v2_coco_2017_11_17' 34 | }, 35 | 'faster_rcnn_inception_v2_coco': { 36 | 'config': 'faster_rcnn_inception_v2_coco.config', 37 | 'url': 'http://download.tensorflow.org/models/object_detection/faster_rcnn_inception_v2_coco_2018_01_28.tar.gz', 38 | 'checkpoint': 'faster_rcnn_inception_v2_coco_2018_01_28' 39 | }, 40 | 'faster_rcnn_resnet50_coco': { 41 | 'config': 'faster_rcnn_resnet50_coco.config', 42 | 'url': 'http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet50_coco_2018_01_28.tar.gz', 43 | 'checkpoint': 'faster_rcnn_resnet50_coco_2018_01_28' 44 | }, 45 | 'rfcn_resnet101_coco': { 46 | 'config': 'rfcn_resnet101_coco.config', 47 | 'url': 'http://download.tensorflow.org/models/object_detection/rfcn_resnet101_coco_2018_01_28.tar.gz', 48 | 'checkpoint': 'rfcn_resnet101_coco_2018_01_28' 49 | }, 50 | 'faster_rcnn_resnet101_coco': { 51 | 'config': 'faster_rcnn_resnet101_coco.config', 52 | 'url': 'http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet101_coco_2018_01_28.tar.gz', 53 | 'checkpoint': 'faster_rcnn_resnet101_coco_2018_01_28' 54 | }, 55 | 'faster_rcnn_inception_resnet_v2_atrous_coco': { 56 | 'config': 'faster_rcnn_inception_resnet_v2_atrous_coco.config', 57 | 'url': 'http://download.tensorflow.org/models/object_detection/faster_rcnn_inception_resnet_v2_atrous_coco_2018_01_28.tar.gz', 58 | 'checkpoint': 'faster_rcnn_inception_resnet_v2_atrous_coco_2018_01_28' 59 | }, 60 | 'faster_rcnn_nas': { 61 | 'config': 'faster_rcnn_nas_coco.config', 62 | 'url': 'http://download.tensorflow.org/models/object_detection/faster_rcnn_nas_coco_2018_01_28.tar.gz', 63 | 'checkpoint': 'faster_rcnn_nas_coco_2018_01_28' 64 | } 65 | } 66 | 67 | def maybe_download_and_extract(url, output_dir): 68 | """Download and extract model tar file. 69 | 70 | If the pretrained model we're using doesn't already exist, this function 71 | downloads it from the TensorFlow.org website and unpacks it into a directory. 72 | """ 73 | if not os.path.exists(output_dir): 74 | os.makedirs(output_dir) 75 | filename = url.split('/')[-1] 76 | filepath = os.path.join(output_dir, filename) 77 | if not os.path.exists(filepath): 78 | def _progress(count, block_size, total_size): 79 | sys.stdout.write('\r>> Downloading %s %.1f%%' % 80 | (filename, 81 | float(count * block_size) / float(total_size) * 100.0)) 82 | sys.stdout.flush() 83 | 84 | filepath, _ = urllib.request.urlretrieve(url, filepath, _progress) 85 | statinfo = os.stat(filepath) 86 | print('Successfully downloaded', filename, statinfo.st_size, 'bytes.') 87 | tarfile.open(filepath, 'r:gz').extractall(output_dir) 88 | 89 | def main(_): 90 | arch_details = arch_map[FLAGS.architecture] 91 | # check graph type, download graph 92 | graph_url = arch_details['url'] 93 | graph_path = '/models/research/object_detection/data/' 94 | maybe_download_and_extract(graph_url, graph_path) 95 | # Open config file 96 | config_path = os.path.join('/models/research/object_detection/samples/configs', 97 | arch_details['config']) 98 | configs = config_util.get_configs_from_pipeline_file(config_path) 99 | # Update paths in config 100 | hparams = tf.contrib.training.HParams(label_map_path=FLAGS.label_map_path, 101 | train_input_path=os.path.join(FLAGS.data_dir, 'train.record'), 102 | eval_input_path=os.path.join(FLAGS.data_dir, 'val.record')) 103 | 104 | if FLAGS.hparams: 105 | for key, val in json.loads(FLAGS.hparams).iteritems(): 106 | hparams.add_hparam(key, val) 107 | 108 | config_util.merge_external_params_with_configs(configs, hparams) 109 | # Save config inside dataset 110 | 111 | configs["train_config"].fine_tune_checkpoint = os.path.join(graph_path, 112 | arch_details['checkpoint'], 'model.ckpt') 113 | 114 | config_proto = config_util.create_pipeline_proto_from_configs(configs) 115 | config_str = text_format.MessageToString(config_proto) 116 | 117 | experiment_path = os.path.join(FLAGS.data_dir, FLAGS.experiment_id) 118 | if not os.path.exists(experiment_path): 119 | os.makedirs(experiment_path) 120 | with open(os.path.join(experiment_path, 'pipeline.config'), 'w') as config_file: 121 | config_file.write(config_str) 122 | 123 | if __name__ == '__main__': 124 | tf.app.run() -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | # Usage info 5 | show_help() { 6 | cat << EOF 7 | Usage: ${0##*/} [-m mode] [-a architecture] [-h hparams] [-e experiment_id] [-c checkpoint] 8 | 9 | -h display this help and exit 10 | -m mode: should be either `train` or `export` 11 | -p key value pairs of hyperparameters as json string 12 | -e experiment id. Used as path inside data folder to run current experiment 13 | -c applicable when mode is export, used to specify checkpoint to use for export 14 | EOF 15 | } 16 | 17 | ARCHITECTURE="ssd_mobilenet_v1_coco" 18 | EXPERIMENT_ID="0" 19 | HPARAMS="" 20 | DATA_DIR="/data" 21 | LABEL_MAP_PATH="/data/label_map.pbtxt" 22 | CHECKPOINT_FILE="model.ckpt" 23 | 24 | MODE="train" 25 | OPTIND=1 26 | 27 | while getopts m:a:h:e:c:p: opt; do 28 | case $opt in 29 | m) MODE=$OPTARG 30 | ;; 31 | a) ARCHITECTURE=$OPTARG 32 | ;; 33 | p) HPARAMS=$OPTARG 34 | ;; 35 | e) EXPERIMENT_ID=$OPTARG 36 | ;; 37 | c) CHECKPOINT_FILE=$OPTARG 38 | ;; 39 | h) 40 | show_help >&2 41 | exit 1 42 | ;; 43 | *) 44 | show_help >&2 45 | exit 1 46 | ;; 47 | esac 48 | done 49 | 50 | echo "MODE: $MODE" 51 | echo "ARCHITECTURE: $ARCHITECTURE" 52 | echo "EXPERIMENT ID: $EXPERIMENT_ID" 53 | echo "HPARAMS: $HPARAMS" 54 | 55 | TRAIN_DIR="$DATA_DIR/$EXPERIMENT_ID" 56 | 57 | if [ $MODE == "train" ] 58 | then 59 | # Create label map file from dataset 60 | python /python/create_label_map.py \ 61 | --data_dir $DATA_DIR \ 62 | --label_map_path $LABEL_MAP_PATH 63 | 64 | # Create tf records from dataset 65 | python /python/create_data_tf_record.py \ 66 | --data_dir $DATA_DIR \ 67 | --output_dir $DATA_DIR \ 68 | --label_map_path $LABEL_MAP_PATH 69 | 70 | if [ ! -z "$HPARAMS" -a "$HPARAMS" != " " ]; then 71 | # Create config file 72 | python /python/update_config.py \ 73 | --architecture $ARCHITECTURE \ 74 | --experiment_id $EXPERIMENT_ID \ 75 | --label_map_path $LABEL_MAP_PATH \ 76 | --data_dir $DATA_DIR \ 77 | --hparams $HPARAMS 78 | else 79 | # Create config file 80 | python /python/update_config.py \ 81 | --architecture $ARCHITECTURE \ 82 | --experiment_id $EXPERIMENT_ID \ 83 | --label_map_path $LABEL_MAP_PATH \ 84 | --data_dir $DATA_DIR 85 | fi 86 | 87 | mkdir -p "$TRAIN_DIR/eval" 88 | 89 | # Start eval on cpu 90 | nohup bash -c "sleep 30; 91 | env CUDA_VISIBLE_DEVICES=-1 python /models/research/object_detection/eval.py \ 92 | --checkpoint_dir $TRAIN_DIR \ 93 | --eval_dir \"$TRAIN_DIR/eval\" \ 94 | --pipeline_config_path \"$TRAIN_DIR/pipeline.config\"" & 95 | 96 | # Start tensorboard at port 8000 97 | nohup tensorboard --port 8000 --logdir=$TRAIN_DIR & 98 | 99 | # Start training 100 | python /models/research/object_detection/train.py \ 101 | --train_dir $TRAIN_DIR \ 102 | --pipeline_config_path "$TRAIN_DIR/pipeline.config" 103 | 104 | elif [ $MODE = "export" ] 105 | then 106 | # Export last trained model in experiment 107 | python /models/research/object_detection/export_inference_graph.py \ 108 | --trained_checkpoint_prefix $CHECKPOINT_FILE \ 109 | --output_directory $TRAIN_DIR \ 110 | --pipeline_config_path "$TRAIN_DIR/pipeline.config" 111 | 112 | /tensorflow/bazel-bin/tensorflow/tools/graph_transforms/transform_graph \ 113 | --in_graph="$TRAIN_DIR/frozen_inference_graph.pb" \ 114 | --out_graph="$TRAIN_DIR/quantized_graph.pb" \ 115 | --inputs='image_tensor' \ 116 | --outputs='detection_boxes,detection_scores,detection_classes,num_detections' \ 117 | --transforms='quantize_weights' 118 | fi --------------------------------------------------------------------------------