Keep dots in sample names for YOLO and COCO datasets

perceptionmetrics/datasets/coco.py

@@ -92,7 +92,8 @@ def build_coco_dataset(
     dataset = OrderedDict()
     for img_id in coco.getImgIds():
         img_info = coco.loadImgs(img_id)[0]
-        sample_name = os.path.basename(img_info["file_name"]).split(".")[0]
+        # splitext keeps dots in the stem, see yolo.build_dataset
+        sample_name = os.path.splitext(os.path.basename(img_info["file_name"]))[0]
         dataset[sample_name] = (
             img_info["file_name"],
             str(img_id),

perceptionmetrics/datasets/yolo.py

@@ -69,7 +69,10 @@ def _make_path_abs(p: str) -> str:
         def _add_to_dataset(image_fname: str, label_fname: str, split: str) -> None:
             """Add a sample to the dataset DataFrame"""
             if os.path.isfile(image_fname) and os.path.isfile(label_fname):
-                sample_name = os.path.basename(image_fname).split(".")[0]
+                # splitext, not split("."): Roboflow exports name files like
+                # "IMG_0001_jpg.rf.<hash>.jpg", and cutting at the first dot gives
+                # every augmented copy the same key, so all but one get dropped
+                sample_name = os.path.splitext(os.path.basename(image_fname))[0]
                 dataset[sample_name] = (
                     os.path.relpath(image_fname, dataset_dir),
                     os.path.relpath(label_fname, dataset_dir),

tests/datasets/test_coco.py

@@ -0,0 +1,64 @@
+import json
+
+from perceptionmetrics.datasets.coco import build_coco_dataset
+
+
+def _write_coco(tmp_path, file_names):
+    """Write a minimal COCO annotation file with one box per image.
+
+    :param tmp_path: Directory to write into
+    :type tmp_path: pathlib.Path
+    :param file_names: Image file names to list in the annotation file
+    :type file_names: List[str]
+    :return: Path to the annotation file and to the image directory
+    :rtype: Tuple[str, str]
+    """
+    image_dir = tmp_path / "images"
+    image_dir.mkdir()
+    images, annotations = [], []
+    for i, name in enumerate(file_names, start=1):
+        images.append({"id": i, "file_name": name, "width": 64, "height": 48})
+        annotations.append(
+            {
+                "id": i,
+                "image_id": i,
+                "category_id": 1,
+                "bbox": [8, 8, 16, 16],
+                "area": 256,
+                "iscrowd": 0,
+            }
+        )
+    ann_file = tmp_path / "instances_val.json"
+    ann_file.write_text(
+        json.dumps(
+            {
+                "images": images,
+                "annotations": annotations,
+                "categories": [{"id": 1, "name": "helmet"}],
+            }
+        )
+    )
+    return str(ann_file), str(image_dir)
+
+
+def test_build_coco_dataset_keeps_samples_with_dots_in_name(tmp_path):
+    """Every image in the annotation file must get its own row, even when the
+    file names only differ after the first dot (e.g. Roboflow exports).
+
+    :param tmp_path: pytest temporary directory
+    :type tmp_path: pathlib.Path
+    """
+    names = [
+        "IMG_0001_jpg.rf.3f9a1c7e2b.jpg",
+        "IMG_0001_jpg.rf.8b21d04aa9.jpg",
+        "IMG_0002_jpg.rf.1a2b3c4d5e.jpg",
+        "plain.jpg",
+    ]
+    ann_file, image_dir = _write_coco(tmp_path, names)
+
+    dataset, _ = build_coco_dataset(ann_file, image_dir, split="val")
+
+    assert len(dataset) == 4
+    assert list(dataset["annotation"]) == ["1", "2", "3", "4"]
+    assert "IMG_0001_jpg.rf.8b21d04aa9" in dataset.index
+    assert "plain" in dataset.index

tests/datasets/test_yolo.py

@@ -136,3 +136,30 @@ def test_build_dataset(caplog: pytest.LogCaptureFixture) -> None:
     assert "cat" in ontology and "dog" in ontology
     assert ontology["cat"]["idx"] == 0
     assert ontology["dog"]["idx"] == 1
+
+
+def test_build_dataset_keeps_samples_with_dots_in_name() -> None:
+    """Samples whose file names contain extra dots must not collapse into one.
+
+    Roboflow exports name every augmented copy of an image
+    ``<original>_jpg.rf.<hash>.jpg``, so the part before the first dot is the
+    same for all of them.
+    """
+    labels = [
+        "/fake/dataset/labels/val/IMG_0001_jpg.rf.3f9a1c7e2b.txt",
+        "/fake/dataset/labels/val/IMG_0001_jpg.rf.8b21d04aa9.txt",
+        "/fake/dataset/labels/val/IMG_0002_jpg.rf.1a2b3c4d5e.txt",
+    ]
+    dataset, _, _ = _make_patched_build_dataset(
+        _FAKE_YAML_TRAIN_VAL_ONLY, {"val": labels}
+    )
+
+    assert len(dataset) == 3
+    assert "IMG_0001_jpg.rf.3f9a1c7e2b" in dataset.index
+    assert "IMG_0001_jpg.rf.8b21d04aa9" in dataset.index
+
+    # Plain file names keep the same sample name as before
+    dataset, _, _ = _make_patched_build_dataset(
+        _FAKE_YAML_TRAIN_VAL_ONLY, {"val": ["/fake/dataset/labels/val/img1.txt"]}
+    )
+    assert list(dataset.index) == ["img1"]