diff --git a/src/labelformat/formats/kitti.py b/src/labelformat/formats/kitti.py index 400fbbb..d2e0154 100644 --- a/src/labelformat/formats/kitti.py +++ b/src/labelformat/formats/kitti.py @@ -126,9 +126,10 @@ def save(self, label_input: ObjectDetectionInput) -> None: left, top, right, bottom = obj.box.to_format( format=BoundingBoxFormat.XYXY ) + category_name = obj.category.name.replace(" ", "_") # Unknown values match Kitti dataset "DontCare" label values. file.write( - f"{obj.category.name} " + f"{category_name} " "-1 " # truncated "-1 " # occluded "-10 " # alpha diff --git a/tests/unit/formats/test_kitti.py b/tests/unit/formats/test_kitti.py index 6cd7cfe..5a39840 100644 --- a/tests/unit/formats/test_kitti.py +++ b/tests/unit/formats/test_kitti.py @@ -1,6 +1,5 @@ from argparse import ArgumentParser from pathlib import Path -from typing import Iterable import pytest from pytest_mock import MockerFixture @@ -9,12 +8,12 @@ KittiObjectDetectionInput, KittiObjectDetectionOutput, ) +from labelformat.formats.labelformat import LabelformatObjectDetectionInput from labelformat.model.bounding_box import BoundingBox from labelformat.model.category import Category from labelformat.model.image import Image from labelformat.model.object_detection import ( ImageObjectDetection, - ObjectDetectionInput, SingleObjectDetection, ) @@ -96,3 +95,32 @@ def test_save(self, tmp_path: Path, with_confidence: bool) -> None: "cat -1 -1 -10 50.0 60.0 70.0 80.0 -1 -1 -1 -1000 -1000 -1000 -10\n" ) assert contents == expected + + def test_save_replaces_spaces_in_category_name(self, tmp_path: Path) -> None: + output_folder = tmp_path / "labels" + categories = [ + Category(id=0, name="traffic light"), + ] + image = Image(id=0, filename="image.jpg", width=100, height=200) + label_input = LabelformatObjectDetectionInput( + categories=categories, + images=[image], + labels=[ + ImageObjectDetection( + image=image, + objects=[ + SingleObjectDetection( + category=categories[0], + box=BoundingBox(xmin=10.0, ymin=20.0, xmax=30.0, ymax=40.0), + ), + ], + ) + ], + ) + + KittiObjectDetectionOutput(output_folder=output_folder).save( + label_input=label_input + ) + + contents = (output_folder / "image.txt").read_text() + assert contents.startswith("traffic_light ")