Factor out the images.meta loading code from YoloExtractor

datumaro/plugins/yolo_format/extractor.py

@@ -1,5 +1,5 @@
 
-# Copyright (C) 2019-2020 Intel Corporation
+# Copyright (C) 2019-2021 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 
@@ -11,7 +11,9 @@
     AnnotationType, Bbox, DatasetItem, Extractor, Importer, LabelCategories,
     SourceExtractor,
 )
-from datumaro.util.image import Image
+from datumaro.util.image import (
+    DEFAULT_IMAGE_META_FILE_NAME, Image, load_image_meta_file,
+)
 from datumaro.util.os_util import split_path
 
 from .format import YoloPath
@@ -47,17 +49,12 @@ def __init__(self, config_path, image_info=None):
 
         assert image_info is None or isinstance(image_info, (str, dict))
         if image_info is None:
-            image_info = osp.join(rootpath, YoloPath.IMAGE_META_FILE)
+            image_info = osp.join(rootpath, DEFAULT_IMAGE_META_FILE_NAME)
             if not osp.isfile(image_info):
                 image_info = {}
         if isinstance(image_info, str):
-            if not osp.isfile(image_info):
-                raise Exception("Can't read image meta file '%s'" % image_info)
-            with open(image_info, encoding='utf-8') as f:
-                image_info = {}
-                for line in f:
-                    image_name, h, w = line.strip().rsplit(maxsplit=2)
-                    image_info[image_name] = (int(h), int(w))
+            image_info = load_image_meta_file(image_info)
+
         self._image_info = image_info
 
         with open(config_path, 'r', encoding='utf-8') as f:

datumaro/plugins/yolo_format/format.py

@@ -1,11 +1,9 @@
 
-# Copyright (C) 2019-2020 Intel Corporation
+# Copyright (C) 2019-2021 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 
 
 class YoloPath:
     DEFAULT_SUBSET_NAME = 'train'
     SUBSET_NAMES = ['train', 'valid']
-
-    IMAGE_META_FILE = 'images.meta'

datumaro/util/image.py

@@ -4,10 +4,13 @@
 
 from enum import Enum, auto
 from io import BytesIO
-from typing import Any, Callable, Iterable, Iterator, Optional, Tuple, Union
+from typing import (
+    Any, Callable, Dict, Iterable, Iterator, Optional, Tuple, Union,
+)
 import importlib
 import os
 import os.path as osp
+import shlex
 import shutil
 
 import numpy as np
@@ -380,3 +383,39 @@ def save(self, path):
                 f.write(self.get_bytes())
         else:
             save_image(path, self.data)
+
+ImageMeta = Dict[str, Tuple[int, int]]
+
+DEFAULT_IMAGE_META_FILE_NAME = 'images.meta'
+
+def load_image_meta_file(image_meta_path: str) -> ImageMeta:
+    """
+    Loads image metadata from a file with the following format:
+
+        <image name 1> <height 1> <width 1>
+        <image name 2> <height 2> <width 2>
+        ...
+
+    Shell-like comments and quoted fields are allowed.
+
+    This can be useful to support datasets in which image dimensions are
+    required to interpret annotations.
+    """
+    assert isinstance(image_meta_path, str)
+
+    if not osp.isfile(image_meta_path):
+        raise Exception("Can't read image meta file '%s'" % image_meta_path)
+
+    image_meta = {}
+
+    with open(image_meta_path, encoding='utf-8') as f:
+        for line in f:
+            fields = shlex.split(line, comments=True)
+            if not fields:
+                continue
+
+            # ignore extra fields, so that the format can be extended later
+            image_name, h, w = fields[:3]
+            image_meta[image_name] = (int(h), int(w))
+
+    return image_meta

tests/test_images.py

@@ -4,7 +4,8 @@
 import numpy as np
 
 from datumaro.util.image import (
-    ByteImage, Image, encode_image, lazy_image, load_image, save_image,
+    ByteImage, Image, encode_image, lazy_image, load_image,
+    load_image_meta_file, save_image,
 )
 from datumaro.util.image_cache import ImageCache
 from datumaro.util.test_utils import TestDir
@@ -130,3 +131,28 @@ def test_ctors(self):
                     if 'ext' in args or 'path' in args:
                         self.assertEqual(img.ext, args.get('ext', '.png'))
                     # pylint: enable=pointless-statement
+
+class ImageMetaTest(TestCase):
+    @mark_requirement(Requirements.DATUM_GENERAL_REQ)
+    def test_loading(self):
+        meta_file_contents = r"""
+        # this is a comment
+
+        a 123 456
+        'b c' 10 20 # inline comment
+        """
+
+        meta_expected = {
+            'a': (123, 456),
+            'b c': (10, 20),
+        }
+
+        with TestDir() as test_dir:
+            meta_path = osp.join(test_dir, 'images.meta')
+
+            with open(meta_path, 'w') as meta_file:
+                meta_file.write(meta_file_contents)
+
+            meta_loaded = load_image_meta_file(meta_path)
+
+        self.assertEqual(meta_loaded, meta_expected)