diff --git a/doctr/datasets/recognition.py b/doctr/datasets/recognition.py index 714be882ba..15011251b6 100644 --- a/doctr/datasets/recognition.py +++ b/doctr/datasets/recognition.py @@ -47,10 +47,10 @@ def __init__( self.data.append((img_name, label)) def merge_dataset(self, ds: AbstractDataset) -> None: - # Update data with new root for self - self.data = [(str(Path(self.root).joinpath(img_path)), label) for img_path, label in self.data] + # Update data with new root for self (absolute paths, a relative one would be read from the new root) + self.data = [(str(Path(self.root).absolute().joinpath(img_path)), label) for img_path, label in self.data] # Define new root self.root = Path("/") # Merge with ds data for img_path, label in ds.data: - self.data.append((str(Path(ds.root).joinpath(img_path)), label)) + self.data.append((str(Path(ds.root).absolute().joinpath(img_path)), label)) diff --git a/tests/pytorch/test_datasets_pt.py b/tests/pytorch/test_datasets_pt.py index 7c0030ea00..a5040cb1fe 100644 --- a/tests/pytorch/test_datasets_pt.py +++ b/tests/pytorch/test_datasets_pt.py @@ -350,7 +350,7 @@ def test_table_dataset(mock_image_folder, mock_table_label, use_polygons): move(os.path.join(mock_image_folder, "tmp_file"), os.path.join(mock_image_folder, img_name)) -def test_recognition_dataset(mock_image_folder, mock_recognition_label): +def test_recognition_dataset(mock_image_folder, mock_recognition_label, monkeypatch): input_size = (32, 128) ds = datasets.RecognitionDataset( img_folder=mock_image_folder, @@ -377,6 +377,26 @@ def test_recognition_dataset(mock_image_folder, mock_recognition_label): datasets.RecognitionDataset(mock_image_folder, mock_recognition_label) move(os.path.join(ds.root, "tmp_file"), os.path.join(ds.root, img_name)) + # Merge datasets with a relative image folder + monkeypatch.chdir(os.path.dirname(mock_image_folder)) + img_folder = os.path.basename(mock_image_folder) + ds = datasets.RecognitionDataset( + img_folder=img_folder, + labels_path=mock_recognition_label, + img_transforms=Resize(input_size, preserve_aspect_ratio=True), + ) + labels = [label for _, label in ds.data] + ds.merge_dataset(datasets.RecognitionDataset(img_folder, mock_recognition_label)) + ds.merge_dataset(datasets.RecognitionDataset(img_folder, mock_recognition_label)) + assert len(ds) == 15 + + loader = DataLoader(ds, batch_size=5, collate_fn=ds.collate_fn) + merged_labels = [] + for images, batch_labels in loader: + assert isinstance(images, torch.Tensor) and images.shape == (5, 3, *input_size) + merged_labels.extend(batch_labels) + assert merged_labels == 3 * labels + @pytest.mark.parametrize( "use_polygons",