Fix issue with incorrect image filename generated in annotations; add…

…ed handling for home dir shortcut. (#296)

tools/cloud-vision-utils/cloud_vision_utils/annotation.py

@@ -83,30 +83,35 @@ def read_from_pascal(filename: str):
   Args:
     filename: PASCAL VOC XML filename.
 
-  Yields:
-    `BoundingBox` objects.
+  Returns:
+    (image filename, list of `BoundingBox` objects).
   """
 
   with gfile.GFile(filename) as f:
     tree = ET.parse(f)
 
   root = tree.getroot()
+  image_filename = root.find('filename').text
   size = root.find('size')
   width = int(size.find('width').text)
   height = int(size.find('height').text)
 
+  bounding_boxes = []
   for obj in root.iter('object'):
     # Expected one bounding box per object.
     b = obj.find('bndbox')
     if b:
       label = obj.find('name').text or ''
-      yield BoundingBox(
+      bounding_box = BoundingBox(
           _safe_divide(int(b.find('xmin').text), width),
           _safe_divide(int(b.find('ymin').text), height),
           _safe_divide(int(b.find('xmax').text), width),
           _safe_divide(int(b.find('ymax').text), height),
           label,
       )
+      bounding_boxes.append(bounding_box)
     else:
       raise AttributeError('Could not find "bndbox" element')
 
+  return image_filename, bounding_boxes
+

tools/cloud-vision-utils/cloud_vision_utils/annotation_test.py

@@ -39,19 +39,19 @@ def _test_helper(self, exception=None):
     with tempfile.NamedTemporaryFile(suffix='.xml') as f:
       f.write(ET.tostring(self.root))
       f.seek(0)
-      boxes = annotation.read(f.name)
       if exception:
         with self.assertRaises(exception):
-          _ = list(boxes)
+          _ = annotation.read(f.name)
       else:
-        return list(boxes)
+        image_filename, boxes = annotation.read(f.name)
+        return image_filename, boxes
 
   def test_nominal_case(self):
     """Test successfully reading bounding boxes from PASCAL VOC XML file."""
 
-    boxes = list(annotation.read(self.filename))
+    image_filename, boxes = list(annotation.read(self.filename))
+    self.assertEqual(image_filename, 'image.jpg')
     self.assertEqual(len(boxes), 2)
-
     width = 400
     height = 300
     b = boxes[0]
@@ -91,7 +91,7 @@ def test_one_object_multi_bndbox(self):
     bndbox = obj.find('bndbox')
     bndbox.find('xmin').text = text_num
     obj.append(bndbox)
-    boxes = self._test_helper()
+    _, boxes = self._test_helper()
     self.assertEqual(len(boxes), 2)
     self.assertNotEqual(boxes[0].xmin, int(text_num))
 

tools/cloud-vision-utils/cloud_vision_utils/dataset.py

@@ -50,9 +50,9 @@ def gen_csv_from_images(
 
   get_label = basename if add_label else lambda _: ''
 
-  with gfile.GFile(output_file, 'w') as f:
+  with gfile.GFile(os.path.expanduser(output_file), 'w') as f:
     writer = csv.writer(f, delimiter=',')
-    for topdir, _, files in gfile.walk(input_dir):
+    for topdir, _, files in gfile.walk(os.path.expanduser(input_dir)):
       for f in files:
         label = get_label(topdir)
         row = ([dataset_type, f, label] +
@@ -82,15 +82,16 @@ def gen_csv_from_annotations(
   if not gfile.exists(input_dir):
     raise ValueError('Input directory not found.')
 
-  with gfile.GFile(output_file, 'w') as outf:
+  with gfile.GFile(os.path.expanduser(output_file), 'w') as outf:
     writer = csv.writer(outf, delimiter=',')
-    for filename in gfile.listdir(input_dir):
+    for filename in gfile.listdir(os.path.expanduser(input_dir)):
       filepath = os.path.join(input_dir, filename)
-      for b in annotation.read(filepath):
-        filename = os.path.join(out_path_prefix, filename)
+      image_filename, boxes = annotation.read(filepath)
+      out_image_filename = os.path.join(out_path_prefix, image_filename)
+      for b in boxes:
         row = [
             dataset_type,
-            filename,
+            out_image_filename,
             b.label,
             b.xmin,
             b.ymin,

tools/cloud-vision-utils/cloud_vision_utils/dataset_test.py

@@ -25,8 +25,8 @@
 import mock
 import testfixtures
 
-from cloud_vision_utils import dataset
 from cloud_vision_utils import annotation
+from cloud_vision_utils import dataset
 
 
 class GenCsvFromImagesTest(unittest.TestCase):
@@ -94,7 +94,7 @@ def _test_helper(self, *args, **kwargs):
         d.write(f, b'any')
 
       f = tempfile.NamedTemporaryFile(mode='r+t', suffix='.csv')
-      mock_ret_val = [self.bounding_box]*self.NUM_BOUNDING_BOXES
+      mock_ret_val = ('image.csv', [self.bounding_box]*self.NUM_BOUNDING_BOXES)
       with mock.patch.object(
           annotation, 'read', return_value=mock_ret_val):
         dataset.gen_csv_from_annotations(d.path, f.name, *args, **kwargs)