PR draft of new YOLO and COCO data set loaders

data/xView.yaml

@@ -36,7 +36,7 @@ download: |
   from PIL import Image
   from tqdm import tqdm
 
-  from utils.datasets import autosplit
+  from utils.datasets_old import autosplit
   from utils.general import download, xyxy2xywhn
 
 

detect.py

@@ -20,7 +20,7 @@
 sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path
 
 from models.experimental import attempt_load
-from utils.datasets import LoadStreams, LoadImages
+from utils.datasets_old import LoadStreams, LoadImages
 from utils.general import check_img_size, check_requirements, check_imshow, colorstr, is_ascii, non_max_suppression, \
     apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
 from utils.plots import Annotator, colors

models/common.py

@@ -17,7 +17,7 @@
 from PIL import Image
 from torch.cuda import amp
 
-from utils.datasets import exif_transpose, letterbox
+from utils.datasets_old import exif_transpose, letterbox
 from utils.general import colorstr, increment_path, is_ascii, make_divisible, non_max_suppression, save_one_box, \
     scale_coords, xyxy2xywh
 from utils.plots import Annotator, colors

models/tf.py

@@ -52,7 +52,7 @@
 from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, autopad, C3
 from models.experimental import MixConv2d, CrossConv, attempt_load
 from models.yolo import Detect
-from utils.datasets import LoadImages
+from utils.datasets_old import LoadImages
 from utils.general import make_divisible, check_file, check_dataset
 
 logger = logging.getLogger(__name__)

requirements.txt

@@ -30,3 +30,4 @@ pandas
 # pycocotools>=2.0  # COCO mAP
 # albumentations>=1.0.3
 thop  # FLOPs computation
+# pytest

tests/utils/test_file.py

@@ -0,0 +1,62 @@
+import shutil
+from pathlib import Path
+from typing import Generator, Optional, List, Callable
+
+import pytest
+
+from tests.utils.test_utils import prepare_temporary_dir
+from utils.file import dump_text_file, get_directory_content
+
+
+@pytest.fixture
+def mock_directory_path() -> Generator[str, None, None]:
+    output_path = prepare_temporary_dir(directory_name="mock_directory_path")
+    yield output_path
+    shutil.rmtree(output_path)
+
+
+def mock_directory_content(directory_path: str) -> None:
+    dump_text_file(Path(directory_path).joinpath('file_1.json').as_posix(), '')
+    dump_text_file(Path(directory_path).joinpath('file_2.txt').as_posix(), '')
+    dump_text_file(Path(directory_path).joinpath('file_3.txt').as_posix(), '')
+
+
+@pytest.mark.parametrize(
+    "extension, mock_callback, expected_result",
+    [
+        (
+            None,
+            lambda x: None,
+            0
+        ),  # empty directory
+        (
+            None,
+            mock_directory_content,
+            3
+        ),  # directory contain 3 files
+        (
+            'json',
+            mock_directory_content,
+            1
+        ),  # directory contain 1 .json file
+        (
+            'txt',
+            mock_directory_content,
+            2
+        ),  # directory contain 2 .txt files
+        (
+            'avi',
+            mock_directory_content,
+            0
+        ),  # directory contain 0 .avi files
+    ]
+)
+def test_get_directory_content(
+    mock_directory_path: str,
+    extension: Optional[str],
+    mock_callback: Callable[[str], None],
+    expected_result: List[str]
+) -> None:
+    mock_callback(mock_directory_path)
+    result = get_directory_content(directory_path=mock_directory_path, extension=extension)
+    assert len(result) == expected_result

tests/utils/test_utils.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+
+
+def prepare_temporary_dir(directory_name: str) -> str:
+    directory_path = Path(__file__).parent.joinpath(directory_name)
+    directory_path.mkdir(parents=True, exist_ok=True)
+    return directory_path.as_posix()

train.py

@@ -33,7 +33,7 @@
 from models.experimental import attempt_load
 from models.yolo import Model
 from utils.autoanchor import check_anchors
-from utils.datasets import create_dataloader
+from utils.datasets_old import create_dataloader
 from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
     strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
     check_requirements, print_mutation, set_logging, one_cycle, colorstr, methods

utils/autoanchor.py

@@ -109,7 +109,7 @@ def print_results(k):
     if isinstance(dataset, str):  # *.yaml file
         with open(dataset, errors='ignore') as f:
             data_dict = yaml.safe_load(f)  # model dict
-        from utils.datasets import LoadImagesAndLabels
+        from utils.datasets_old import LoadImagesAndLabels
         dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
 
     # Get label wh

utils/datasets/coco.py

@@ -0,0 +1,74 @@
+# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
+"""
+COCO dataset loading utils
+"""
+
+import json
+from typing import Dict, Union, List
+from collections import defaultdict
+
+import torch
+import numpy as np
+
+
+IMAGE_KEY = "images"
+IMAGE_FILE_NAME_KEY = "file_name"
+IMAGE_ID_KEY = "id"
+IMAGE_WIDTH_KEY = "width"
+IMAGE_HEIGHT_KEY = "height"
+ANNOTATION_KEY = "annotations"
+ANNOTATION_IMAGE_ID_KEY = "image_id"
+ANNOTATION_BBOX_KEY = "bbox"
+ANNOTATION_CATEGORY_ID = "category_id"
+
+
+def read_json_file(file_path: str, **kwargs) -> Union[list, dict]:
+    with open(file_path, 'r') as file:
+        return json.load(file, **kwargs)
+
+
+def load_coco_annotations(coco_data: dict) -> Dict[str, torch.Tensor]:
+    coco_image_entries_map = map_coco_image_entries(coco_image_entries=coco_data[IMAGE_KEY])
+    coco_annotation_entries_map = map_coco_annotation_entries(coco_annotation_entries=coco_data[ANNOTATION_KEY])
+    return {
+        coco_image_entries_map[image_id][IMAGE_FILE_NAME_KEY]: process_coco_annotation(
+            coco_annotation_entries=coco_annotation_entries_map[image_id],
+            coco_image_data=coco_image_entries_map[image_id]
+        )
+        for image_id
+        in sorted(coco_image_entries_map.keys())
+    }
+
+
+def map_coco_image_entries(coco_image_entries: List[dict]) -> Dict[int, dict]:
+    return {
+        image_data[IMAGE_ID_KEY]: image_data
+        for image_data
+        in coco_image_entries
+    }
+
+
+def map_coco_annotation_entries(coco_annotation_entries: List[dict]) -> Dict[int, List[dict]]:
+    result = defaultdict(list)
+    for coco_annotation_entry in coco_annotation_entries:
+        image_id = coco_annotation_entry[ANNOTATION_IMAGE_ID_KEY]
+        result[image_id].append(coco_annotation_entry)
+    return result
+
+
+def process_coco_annotation(coco_annotation_entries: List[dict], coco_image_data: dict) -> torch.Tensor:
+    image_width = coco_image_data[IMAGE_WIDTH_KEY]
+    image_height = coco_image_data[IMAGE_HEIGHT_KEY]
+    annotations = []
+    for coco_annotation_entry in coco_annotation_entries:
+        category_id = coco_annotation_entry[ANNOTATION_CATEGORY_ID]
+        x_min, y_min, width, height = coco_annotation_entry[ANNOTATION_BBOX_KEY]
+        annotations.append([
+            0,
+            category_id,
+            (x_min + width / 2) / image_width,
+            (y_min + height / 2) / image_height,
+            width / image_width,
+            height / image_height
+        ])
+    return torch.as_tensor(np.array(annotations))