diff --git a/ppocr/data/multi_scale_sampler.py b/ppocr/data/multi_scale_sampler.py index 4ab38fc4e65..45a1b4cd582 100644 --- a/ppocr/data/multi_scale_sampler.py +++ b/ppocr/data/multi_scale_sampler.py @@ -167,5 +167,14 @@ def iter(self): def set_epoch(self, epoch: int): self.epoch = epoch + if self.ds_width: + self.wh_ratio = self.data_source.wh_ratio + self.wh_ratio_sort = self.data_source.wh_ratio_sort + # Refresh cached batches after ratio_list resampling. + self.batchs_in_one_epoch = self.iter() + self.batchs_in_one_epoch_id = [ + i for i in range(len(self.batchs_in_one_epoch)) + ] + def __len__(self): return self.length diff --git a/ppocr/data/simple_dataset.py b/ppocr/data/simple_dataset.py index 57e9b4378c5..a40171301f8 100644 --- a/ppocr/data/simple_dataset.py +++ b/ppocr/data/simple_dataset.py @@ -461,9 +461,32 @@ def wh_aware(self): self.data_lines = data_line_new self.wh_ratio = np.array(wh_ratio) - self.wh_ratio_sort = np.argsort(self.wh_ratio) + self._refresh_wh_ratio_sort() self.data_idx_order_list = list(range(len(self.data_lines))) + def _refresh_wh_ratio_sort(self): + if self._index_map is not None: + sampled_indices = np.asarray(self._index_map, dtype=np.int64) + sampled_order = np.argsort(self.wh_ratio[sampled_indices]) + self.wh_ratio_sort = sampled_indices[sampled_order] + else: + self.wh_ratio_sort = np.argsort(self.wh_ratio) + + def _ensure_index_map(self): + previous_epoch = self._cached_epoch + super(MultiScaleDataSet, self)._ensure_index_map() + if ( + self.ds_width + and self._index_map is not None + and self._cached_epoch != previous_epoch + ): + self._refresh_wh_ratio_sort() + + def reset_data_lines(self, seed=None, epoch=None): + super(MultiScaleDataSet, self).reset_data_lines(seed=seed, epoch=epoch) + if self.ds_width: + self._refresh_wh_ratio_sort() + def resize_norm_img(self, data, imgW, imgH, padding=True): img = data["image"] h = img.shape[0] @@ -498,6 +521,10 @@ def __getitem__(self, properties): # properties is a tuple, contains (width, height, index) img_height = properties[1] idx = properties[2] + + if self._index_map is not None: + self._ensure_index_map() + if self.ds_width and properties[3] is not None: wh_ratio = properties[3] img_width = img_height * ( @@ -505,11 +532,18 @@ def __getitem__(self, properties): ) file_idx = self.wh_ratio_sort[idx] else: - file_idx = self.data_idx_order_list[idx] img_width = properties[0] wh_ratio = None + if self._index_map is not None: + file_idx = self._index_map[idx] + else: + file_idx = self.data_idx_order_list[idx] + + if self._index_map is not None: + data_line = self._all_lines[file_idx] + else: + data_line = self.data_lines[file_idx] - data_line = self.data_lines[file_idx] try: data_line = data_line.decode("utf-8") substr = data_line.strip("\n").split(self.delimiter) diff --git a/tests/ppocr/test_simple_dataset.py b/tests/ppocr/test_simple_dataset.py new file mode 100644 index 00000000000..bed5415565a --- /dev/null +++ b/tests/ppocr/test_simple_dataset.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging + +import cv2 +import numpy as np +import pytest + +from ppocr.data.multi_scale_sampler import MultiScaleSampler +from ppocr.data.simple_dataset import MultiScaleDataSet + + +def _create_dataset_file(tmp_path, prefix, num_samples, width, height): + label_file = tmp_path / f"{prefix}.txt" + + with open(label_file, "w", encoding="utf-8") as f: + for idx in range(num_samples): + image_name = f"{prefix}_{idx}.jpg" + image_path = tmp_path / image_name + + image = np.zeros((height, width, 3), dtype=np.uint8) + cv2.imwrite(str(image_path), image) + + f.write(f"{image_name}\t{prefix}_{idx}\t{width}\t{height}\n") + + return str(label_file) + + +def _build_config( + tmp_path, + label_file_list, + ratio_list, + ds_width, + transforms, + shuffle, +): + return { + "Global": {}, + "Train": { + "dataset": { + "name": "MultiScaleDataSet", + "data_dir": str(tmp_path), + "label_file_list": label_file_list, + "ratio_list": ratio_list, + "ds_width": ds_width, + "transforms": transforms, + }, + "loader": { + "shuffle": shuffle, + }, + }, + } + + +@pytest.mark.parametrize("ds_width", [False, True]) +def test_multiscale_dataset_respects_ratio_sampling(tmp_path, ds_width): + # Dataset A has a smaller aspect ratio than dataset B. When ds_width + # is enabled, sorting the full dataset instead of the sampled subset + # therefore exposes the indexing bug deterministically. + a_file = _create_dataset_file( + tmp_path, + prefix="A", + num_samples=2, + width=10, + height=10, + ) + b_file = _create_dataset_file( + tmp_path, + prefix="B", + num_samples=2, + width=20, + height=10, + ) + + config = _build_config( + tmp_path=tmp_path, + label_file_list=[a_file, b_file], + ratio_list=[0.5, 0.5], + ds_width=ds_width, + transforms=[ + { + "DecodeImage": { + "img_mode": "BGR", + "channel_first": False, + } + }, + { + "KeepKeys": { + "keep_keys": [ + "image", + "label", + ] + } + }, + ], + shuffle=True, + ) + + dataset = MultiScaleDataSet( + config=config, + mode="Train", + logger=logging.getLogger(__name__), + seed=0, + ) + + assert len(dataset) == 2 + + # ratio_list should select one sample from each label file. + sampled_indices = dataset._index_map + assert sum(index < 2 for index in sampled_indices) == 1 + assert sum(index >= 2 for index in sampled_indices) == 1 + + samples = [dataset.__getitem__([320, 48, idx, 1.0]) for idx in range(len(dataset))] + labels = [sample[1] for sample in samples] + + assert sum(label.startswith("A_") for label in labels) == 1 + assert sum(label.startswith("B_") for label in labels) == 1 + + +def test_multiscale_sampler_refreshes_ratio_after_dataset_reset(tmp_path): + label_file = tmp_path / "data.txt" + + sizes = [ + (10, 10), + (20, 10), + (30, 10), + (40, 10), + ] + + with open(label_file, "w", encoding="utf-8") as f: + for idx, (width, height) in enumerate(sizes): + f.write(f"image_{idx}.jpg\tlabel_{idx}\t{width}\t{height}\n") + + config = _build_config( + tmp_path=tmp_path, + label_file_list=[str(label_file)], + ratio_list=[0.5], + ds_width=True, + transforms=[], + shuffle=False, + ) + + dataset = MultiScaleDataSet( + config=config, + mode="Train", + logger=logging.getLogger(__name__), + seed=0, + ) + + sampler = MultiScaleSampler( + data_source=dataset, + scales=[[320, 48]], + first_bs=2, + fix_bs=True, + is_training=True, + seed=0, + ) + + first_expected_ratio = np.mean(dataset.wh_ratio[dataset.wh_ratio_sort]) + + first_epoch_batch = list(iter(sampler))[0] + first_epoch_ratio = first_epoch_batch[0][3] + + assert np.isclose( + first_epoch_ratio, + first_expected_ratio, + ) + + # Find another sampling seed that changes the expected aspect ratio. + for seed in range(1, 100): + dataset.reset_data_lines(seed=seed, epoch=seed) + + second_expected_ratio = np.mean(dataset.wh_ratio[dataset.wh_ratio_sort]) + + if not np.isclose( + first_expected_ratio, + second_expected_ratio, + ): + break + else: + pytest.fail("Could not generate a different sampled subset") + + # program.py calls set_epoch() after reset_data_lines(). + sampler.set_epoch(0) + + second_epoch_batch = list(iter(sampler))[0] + second_epoch_ratio = second_epoch_batch[0][3] + + assert np.isclose( + second_epoch_ratio, + second_expected_ratio, + )