diff --git a/perceptionmetrics/datasets/coco.py b/perceptionmetrics/datasets/coco.py index c55c40ed..19744fc2 100644 --- a/perceptionmetrics/datasets/coco.py +++ b/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), diff --git a/perceptionmetrics/datasets/yolo.py b/perceptionmetrics/datasets/yolo.py index f329ba91..e6d54679 100644 --- a/perceptionmetrics/datasets/yolo.py +++ b/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..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), diff --git a/tests/datasets/test_coco.py b/tests/datasets/test_coco.py new file mode 100644 index 00000000..c21036f0 --- /dev/null +++ b/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 diff --git a/tests/datasets/test_yolo.py b/tests/datasets/test_yolo.py index 6fff5ab8..17c116b7 100644 --- a/tests/datasets/test_yolo.py +++ b/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 + ``_jpg.rf..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"]