Update create_pascal_tf_record.py and create_pet_tf_record.py

PiperOrigin-RevId: 192494267
vivalaakam · Apr 13, 2018 · f98f000 · f98f000
1 parent 2151c44
commit f98f000
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 30 deletions.
diff --git a/research/object_detection/dataset_tools/create_pascal_tf_record.py b/research/object_detection/dataset_tools/create_pascal_tf_record.py
@@ -104,7 +104,7 @@ def dict_to_tf_example(data,
   truncated = []
   poses = []
   difficult_obj = []
-  if data.has_key('object'):
+  if 'object' in data:
     for obj in data['object']:
       difficult = bool(int(obj['difficult']))
       if ignore_difficult_instances and difficult:

diff --git a/research/object_detection/dataset_tools/create_pet_tf_record.py b/research/object_detection/dataset_tools/create_pet_tf_record.py
@@ -136,35 +136,36 @@ def dict_to_tf_example(data,
   poses = []
   difficult_obj = []
   masks = []
-  for obj in data['object']:
-    difficult = bool(int(obj['difficult']))
-    if ignore_difficult_instances and difficult:
-      continue
-    difficult_obj.append(int(difficult))
-
-    if faces_only:
-      xmin = float(obj['bndbox']['xmin'])
-      xmax = float(obj['bndbox']['xmax'])
-      ymin = float(obj['bndbox']['ymin'])
-      ymax = float(obj['bndbox']['ymax'])
-    else:
-      xmin = float(np.min(nonzero_x_indices))
-      xmax = float(np.max(nonzero_x_indices))
-      ymin = float(np.min(nonzero_y_indices))
-      ymax = float(np.max(nonzero_y_indices))
-
-    xmins.append(xmin / width)
-    ymins.append(ymin / height)
-    xmaxs.append(xmax / width)
-    ymaxs.append(ymax / height)
-    class_name = get_class_name_from_filename(data['filename'])
-    classes_text.append(class_name.encode('utf8'))
-    classes.append(label_map_dict[class_name])
-    truncated.append(int(obj['truncated']))
-    poses.append(obj['pose'].encode('utf8'))
-    if not faces_only:
-      mask_remapped = (mask_np != 2).astype(np.uint8)
-      masks.append(mask_remapped)
+  if 'object' in data:
+    for obj in data['object']:
+      difficult = bool(int(obj['difficult']))
+      if ignore_difficult_instances and difficult:
+        continue
+      difficult_obj.append(int(difficult))
+
+      if faces_only:
+        xmin = float(obj['bndbox']['xmin'])
+        xmax = float(obj['bndbox']['xmax'])
+        ymin = float(obj['bndbox']['ymin'])
+        ymax = float(obj['bndbox']['ymax'])
+      else:
+        xmin = float(np.min(nonzero_x_indices))
+        xmax = float(np.max(nonzero_x_indices))
+        ymin = float(np.min(nonzero_y_indices))
+        ymax = float(np.max(nonzero_y_indices))
+
+      xmins.append(xmin / width)
+      ymins.append(ymin / height)
+      xmaxs.append(xmax / width)
+      ymaxs.append(ymax / height)
+      class_name = get_class_name_from_filename(data['filename'])
+      classes_text.append(class_name.encode('utf8'))
+      classes.append(label_map_dict[class_name])
+      truncated.append(int(obj['truncated']))
+      poses.append(obj['pose'].encode('utf8'))
+      if not faces_only:
+        mask_remapped = (mask_np != 2).astype(np.uint8)
+        masks.append(mask_remapped)
 
   feature_dict = {
       'image/height': dataset_util.int64_feature(height),