1# Copyright 2019 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"""Preprocesses COCO minival data for Object Detection evaluation using mean Average Precision.
16
17The 2014 validation images & annotations can be downloaded from:
18http://cocodataset.org/#download
19The minival image ID allowlist, a subset of the 2014 validation set, can be
20found here:
21https://github.com/tensorflow/models/blob/master/research/object_detection/data/mscoco_minival_ids.txt.
22
23This script takes in the original images folder, instances JSON file and
24image ID allowlist and produces the following in the specified output folder:
25A subfolder for allowlisted images (images/), and a file (ground_truth.pbtxt)
26containing an instance of tflite::evaluation::ObjectDetectionGroundTruth.
27"""
28
29import argparse
30import ast
31import collections
32import os
33import shutil
34import sys
35
36from absl import logging
37from tensorflow.lite.tools.evaluation.proto import evaluation_stages_pb2
38
39
40def _get_ground_truth_detections(instances_file,
41                                 allowlist_file=None,
42                                 num_images=None):
43  """Processes the annotations JSON file and returns ground truth data corresponding to allowlisted image IDs.
44
45  Args:
46    instances_file: COCO instances JSON file, usually named as
47      instances_val20xx.json.
48    allowlist_file: File containing COCO minival image IDs to allowlist for
49      evaluation, one per line.
50    num_images: Number of allowlisted images to pre-process. First num_images
51      are chosen based on sorted list of filenames. If None, all allowlisted
52      files are preprocessed.
53
54  Returns:
55    A dict mapping image id (int) to a per-image dict that contains:
56      'filename', 'image' & 'height' mapped to filename & image dimensions
57      respectively
58      AND
59      'detections' to a list of detection dicts, with each mapping:
60        'category_id' to COCO category id (starting with 1) &
61        'bbox' to a list of dimension-normalized [top, left, bottom, right]
62        bounding-box values.
63  """
64  # Read JSON data into a dict.
65  with open(instances_file, 'r') as annotation_dump:
66    data_dict = ast.literal_eval(annotation_dump.readline())
67
68  image_data = collections.OrderedDict()
69
70  # Read allowlist.
71  if allowlist_file is not None:
72    with open(allowlist_file, 'r') as allowlist:
73      image_id_allowlist = set([int(x) for x in allowlist.readlines()])
74  else:
75    image_id_allowlist = [image['id'] for image in data_dict['images']]
76
77  # Get image names and dimensions.
78  for image_dict in data_dict['images']:
79    image_id = image_dict['id']
80    if image_id not in image_id_allowlist:
81      continue
82    image_data_dict = {}
83    image_data_dict['id'] = image_dict['id']
84    image_data_dict['file_name'] = image_dict['file_name']
85    image_data_dict['height'] = image_dict['height']
86    image_data_dict['width'] = image_dict['width']
87    image_data_dict['detections'] = []
88    image_data[image_id] = image_data_dict
89
90  shared_image_ids = set()
91  for annotation_dict in data_dict['annotations']:
92    image_id = annotation_dict['image_id']
93    if image_id in image_data:
94      shared_image_ids.add(image_id)
95
96  output_image_ids = sorted(shared_image_ids)
97  if num_images:
98    if num_images <= 0:
99      logging.warning(
100          '--num_images is %d, hence outputing all annotated images.',
101          num_images)
102    elif num_images > len(shared_image_ids):
103      logging.warning(
104          '--num_images (%d) is larger than the number of annotated images.',
105          num_images)
106    else:
107      output_image_ids = output_image_ids[:num_images]
108
109  for image_id in list(image_data):
110    if image_id not in output_image_ids:
111      del image_data[image_id]
112
113  # Get detected object annotations per image.
114  for annotation_dict in data_dict['annotations']:
115    image_id = annotation_dict['image_id']
116    if image_id not in output_image_ids:
117      continue
118
119    image_data_dict = image_data[image_id]
120    bbox = annotation_dict['bbox']
121    # bbox format is [x, y, width, height]
122    # Refer: http://cocodataset.org/#format-data
123    top = bbox[1]
124    left = bbox[0]
125    bottom = top + bbox[3]
126    right = left + bbox[2]
127    if (top > image_data_dict['height'] or left > image_data_dict['width'] or
128        bottom > image_data_dict['height'] or right > image_data_dict['width']):
129      continue
130    object_d = {}
131    object_d['bbox'] = [
132        top / image_data_dict['height'], left / image_data_dict['width'],
133        bottom / image_data_dict['height'], right / image_data_dict['width']
134    ]
135    object_d['category_id'] = annotation_dict['category_id']
136    image_data_dict['detections'].append(object_d)
137
138  return image_data
139
140
141def _dump_data(ground_truth_detections, images_folder_path, output_folder_path):
142  """Dumps images & data from ground-truth objects into output_folder_path.
143
144  The following are created in output_folder_path:
145    images/: sub-folder for allowlisted validation images.
146    ground_truth.pb: A binary proto file containing all ground-truth
147    object-sets.
148
149  Args:
150    ground_truth_detections: A dict mapping image id to ground truth data.
151      Output of _get_ground_truth_detections.
152    images_folder_path: Validation images folder
153    output_folder_path: folder to output files to.
154  """
155  # Ensure output folders exist.
156  if not os.path.exists(output_folder_path):
157    os.makedirs(output_folder_path)
158  output_images_folder = os.path.join(output_folder_path, 'images')
159  if not os.path.exists(output_images_folder):
160    os.makedirs(output_images_folder)
161  output_proto_file = os.path.join(output_folder_path, 'ground_truth.pb')
162
163  ground_truth_data = evaluation_stages_pb2.ObjectDetectionGroundTruth()
164  for image_dict in ground_truth_detections.values():
165    # Create an ObjectsSet proto for this file's ground truth.
166    detection_result = ground_truth_data.detection_results.add()
167    detection_result.image_id = image_dict['id']
168    detection_result.image_name = image_dict['file_name']
169    for detection_dict in image_dict['detections']:
170      object_instance = detection_result.objects.add()
171      object_instance.bounding_box.normalized_top = detection_dict['bbox'][0]
172      object_instance.bounding_box.normalized_left = detection_dict['bbox'][1]
173      object_instance.bounding_box.normalized_bottom = detection_dict['bbox'][2]
174      object_instance.bounding_box.normalized_right = detection_dict['bbox'][3]
175      object_instance.class_id = detection_dict['category_id']
176    # Copy image.
177    shutil.copy2(
178        os.path.join(images_folder_path, image_dict['file_name']),
179        output_images_folder)
180
181  # Dump proto.
182  with open(output_proto_file, 'wb') as proto_file:
183    proto_file.write(ground_truth_data.SerializeToString())
184
185
186def _parse_args():
187  """Creates a parser that parse the command line arguments.
188
189  Returns:
190    A namespace parsed from command line arguments.
191  """
192  parser = argparse.ArgumentParser(
193      description='preprocess_coco_minival: Preprocess COCO minival dataset')
194  parser.add_argument(
195      '--images_folder',
196      type=str,
197      help='Full path of the validation images folder.',
198      required=True)
199  parser.add_argument(
200      '--instances_file',
201      type=str,
202      help='Full path of the input JSON file, like instances_val20xx.json.',
203      required=True)
204  parser.add_argument(
205      '--allowlist_file',
206      type=str,
207      help='File with COCO image ids to preprocess, one on each line.',
208      required=False)
209  parser.add_argument(
210      '--num_images',
211      type=int,
212      help='Number of allowlisted images to preprocess into the output folder.',
213      required=False)
214  parser.add_argument(
215      '--output_folder',
216      type=str,
217      help='Full path to output images & text proto files into.',
218      required=True)
219  return parser.parse_known_args(args=sys.argv[1:])[0]
220
221
222if __name__ == '__main__':
223  args = _parse_args()
224  ground_truths = _get_ground_truth_detections(args.instances_file,
225                                               args.allowlist_file,
226                                               args.num_images)
227  _dump_data(ground_truths, args.images_folder, args.output_folder)
228