add tests for yolo_detection_format

Author: ChanBong

datumaro/plugins/yolo_detection_format/converter.py

@@ -63,7 +63,7 @@ def _check_dataset(self):
             if subset not in YoloDetectionPath.ALLOWED_SUBSET_NAMES:
                 raise DatasetExportError(
                     f"The allowed subset name is in {YoloDetectionPath.ALLOWED_SUBSET_NAMES}, "
-                    f'so that subset "{subset}" is not allowed.'
+                    f"so that subset '{subset}' is not allowed."
                 )
 
         for must_name in YoloDetectionPath.MUST_SUBSET_NAMES:

datumaro/plugins/yolo_detection_format/extractor.py

@@ -7,6 +7,7 @@
 import yaml
 
 import os.path as osp
+import os
 import re
 from collections import OrderedDict
 from typing import Dict, List, Optional, Tuple, Type, TypeVar, Union, Iterator
@@ -68,6 +69,9 @@ def __init__(
         if not osp.isdir(config_path):
             raise DatasetImportError(f"{config_path} should be a directory.")
 
+        if not osp.isfile(osp.join(config_path, META_FILE)):
+            raise DatasetImportError(f"Can't find {META_FILE} in {config_path}")
+        
         if not urls:
             raise DatasetImportError(
                 f"`urls` should be specified for {self.__class__.__name__}, "
@@ -86,7 +90,32 @@ def __init__(
         self._urls = urls
         self._img_files = self._load_img_files(rootpath)
         self._ann_types = set()
+
+        config = YoloDetectionPath._parse_config(osp.join(config_path, META_FILE))
 
+        subsets = {k: v for k, v in config.items() if k in YoloDetectionPath.ALLOWED_SUBSET_NAMES and v is not None}
+
+        for subset_name, list_path in subsets.items():
+            subset = YoloDetectionExtractor.Subset(subset_name, self)
+
+            if osp.isdir(osp.join(rootpath, list_path)):
+                list_path = osp.join(rootpath, list_path)
+                f = os.listdir(list_path)
+                subset.items = OrderedDict(
+                    (self.name_from_path(p), self.localize_path(p)) for p in f if p.strip()
+                )
+            elif osp.isfile(osp.join(rootpath, list_path)):
+                with open(osp.join(rootpath, list_path), "r", encoding="utf-8") as f:
+                    subset.items = OrderedDict(
+                        (self.name_from_path(p), self.localize_path(p)) for p in f if p.strip()
+                    )
+            else:
+                raise InvalidAnnotationError(f"Can't find '{subset_name}' subset list file")
+            
+            subsets[subset_name] = subset
+
+        self._subsets: Dict[str, YoloDetectionExtractor.Subset] = subsets
+
         self._categories = {
             AnnotationType.label: self._load_categories(
                 osp.join(self._path, META_FILE)
@@ -292,13 +321,11 @@ def _parse_annotations(
 
                 annotations.append(
                     Bbox(
-                        x * image_width,
-                        y * image_height,
-                        w * image_width,
-                        h * image_height,
+                        int(x * image_width),
+                        int(y * image_height),
+                        int(w * image_width),
+                        int(h * image_height),
                         label=label_id,
-                        id=idx,
-                        group=idx,
                     )
                 )
             except Exception as e:

datumaro/plugins/yolo_detection_format/format.py

@@ -5,7 +5,7 @@
 from enum import Enum, auto
 from typing import Dict
 import re
-
+import yaml
 
 class YoloDetectionPath:
     DEFAULT_SUBSET_NAME = "train"
@@ -16,18 +16,10 @@ class YoloDetectionPath:
 
     @staticmethod
     def _parse_config(path: str) -> Dict[str, str]:
-        with open(path, "r", encoding="utf-8") as f:
-            config_lines = f.readlines()
-
-        config = {}
-
-        for line in config_lines:
-            match = re.match(r"^\s*(\w+)\s*=\s*(.+)$", line)
-            if not match:
-                continue
-
-            key = match.group(1)
-            value = match.group(2)
-            config[key] = value
+        with open(path, "r") as fp:
+            loaded = yaml.safe_load(fp.read())
 
-        return config
+        if not isinstance(loaded, dict):
+            raise Exception("Invalid config format")
+        
+        return loaded

datumaro/plugins/yolo_detection_format/importer.py

@@ -53,7 +53,7 @@ def _check_ann_file_impl(cls, fp: TextIOWrapper) -> bool:
             fields = line.rstrip("\n").split(" ")
             if len(fields) != 5:
                 raise DatasetImportError(
-                    f"Yolo format txt file should have 5 fields for each line, "
+                    f"Yolo Detection format txt file should have 5 fields for each line, "
                     f"but the read line has {len(fields)} fields: fields={fields}."
                 )
 
@@ -68,19 +68,11 @@ def _check_ann_file_impl(cls, fp: TextIOWrapper) -> bool:
 
     @classmethod
     def _find_loose(cls, path: str, dirname: str) -> List[Dict[str, Any]]:
-        def _filter_ann_file(fpath: str):
-            try:
-                with open(fpath, "r") as fp:
-                    return cls._check_ann_file_impl(fp)
-            except DatasetImportError:
-                return False
-
         sources = cls._find_sources_recursive(
             path,
             ext=".txt",
             extractor_name="",
             dirname=dirname,
-            file_filter=_filter_ann_file,
             filename="**/*",
             max_depth=1
         )