From ebf05217efae02a6519edb6d24ca037adc0eec17 Mon Sep 17 00:00:00 2001 From: Moenupa <49304833+Moenupa@users.noreply.github.com> Date: Fri, 28 Aug 2026 23:39:16 +0800 Subject: [PATCH 1/2] [feat] add image dump (preview) feature --- pyproject.toml | 1 + src/prep/api/__init__.py | 3 +- src/prep/api/dataio.py | 236 ++++++++++--------------- src/prep/api/formatutils/__init__.py | 3 +- src/prep/api/formatutils/images.py | 45 +++++ src/prep/api/load.py | 156 +++++++++++++++++ src/prep/cli.py | 8 + src/prep/constants.py | 1 + tests/api/formatutils/test_images.py | 28 ++- tests/api/test_dataio.py | 248 ++++++++++++--------------- tests/api/test_load.py | 151 ++++++++++++++++ 11 files changed, 596 insertions(+), 284 deletions(-) create mode 100644 src/prep/api/load.py create mode 100644 tests/api/test_load.py diff --git a/pyproject.toml b/pyproject.toml index 67aa24a..48bc75b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "datasets[vision]<4.8", "typer", "openai", + "filetype", ] [project.urls] diff --git a/src/prep/api/__init__.py b/src/prep/api/__init__.py index cdfbf43..27cf70b 100644 --- a/src/prep/api/__init__.py +++ b/src/prep/api/__init__.py @@ -1,5 +1,6 @@ -from .dataio import OutputActions, PathIO, adaptive_load_dataset +from .dataio import OutputActions, PathIO from .formatter import FormatterPipeline, LoadFn, formatter +from .load import adaptive_load_dataset from .log import get_logger from .transform import ( ImageTransform, diff --git a/src/prep/api/dataio.py b/src/prep/api/dataio.py index fff5460..4b03a96 100644 --- a/src/prep/api/dataio.py +++ b/src/prep/api/dataio.py @@ -1,173 +1,78 @@ -import re +import shutil from collections import defaultdict from dataclasses import dataclass from fnmatch import fnmatch from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +import filetype +import typer from datasets import ( Dataset, - DatasetDict, get_dataset_config_names, get_dataset_split_names, - load_dataset, - load_from_disk, ) from huggingface_hub import HfApi +from PIL import Image as PILImage from rich.filesize import decimal from rich.table import Table -from ..constants import HF_PREFIX, SAVE_PREFIX, SKIP +from ..constants import HF_PREFIX, PREVIEW_PREFIX, SAVE_PREFIX, SKIP, WARN_PREFIX from .formatter import FormatterPipeline, get_registered_pipelines +from .formatutils import iter_images from .log import get_logger -from .types import DataFormat, ProcArgs, Split, get_valid_formats, get_valid_splits +from .types import DataFormat, Split, get_valid_formats, get_valid_splits if TYPE_CHECKING: from collections.abc import Generator -logger = get_logger(__name__) -_FORMATS = { - ".parquet": "parquet", - ".pq": "parquet", - ".json": "json", - ".jsonl": "json", - ".ndjson": "json", - ".csv": "csv", - ".tsv": "csv", - ".arrow": "arrow", - ".txt": "text", - ".text": "text", -} - - -def resolve_remote(source: str) -> tuple[str, str | None, str | None]: - """Resolve a remote dataset source into its components. - - Args: - source (str): dataset ID in the format of `org/data[@subset][:split]`. - - Returns: - tuple[str, str | None, str | None]: A tuple containing the source, subset, and split. - """ - match = re.match( - r"^(?P[^@]+)(?:@(?P[^:]+))?(?::(?P.+))?$", source - ) - if not match: - raise ValueError(f"Invalid source format: {source!r}") - - return ( - match.group("source"), - match.group("subset"), - match.group("split"), - ) - - -def resolve_split(split: str | None, available: list[str]) -> str: - if split is None: - if len(available) == 1: - return available[0] - raise ValueError( - f"Split must be specified for dataset with multiple splits: {available}" - ) +logger = get_logger(__name__) - # direct hit is preferred - if split in available: - return split - # translate if "val" -> "validation" if exists - # we want to be generous here to facilitate IO - if split == "val" and "validation" in available: - return "validation" +def guess_ext(raw: bytes, src_path: str | None) -> str: + ext = filetype.guess_extension(raw) + if ext is None and src_path is not None: + ext = Path(src_path).suffix.lower().lstrip(".") + if ext is None: + raise ValueError("Cannot guess file extension from raw bytes or source path") - raise ValueError(f"Split {split!r} not found in dataset: {available}") + return ext -def load_local(source: str, split: str | None = None) -> Dataset | None: - logger.debug(f"Loading dataset from local disk: {source} (split={split})") +def _write_preview_image(entry: Any, dest: Path) -> Path | None: try: - d = load_from_disk(source) - except Exception: - return None - - if isinstance(d, Dataset): - return d - - if not isinstance(d, DatasetDict): - raise RuntimeError(f"Unexpected dataset type: {type(d)} {d}") - - return d[resolve_split(split, list(map(str, d.keys())))] - - -def load_remote( - source: str, split: str | None = None, args: ProcArgs | None = None -) -> Dataset: - """Load a dataset from Hugging Face Hub. - - Args: - source (str): dataset ID in the format of `org/data[@subset][:split]`. - split (str | None, optional): dataset split name. Defaults to None. - args (ProcArgs | None, optional): processing arguments. Defaults to None. - - Returns: - Dataset: the loaded dataset. - """ - logger.debug(f"Loading dataset from Hugging Face Hub: {source} (split={split})") - source, subset, override_split = resolve_remote(source) - - d = load_dataset( - source, - subset, - split=resolve_split( - override_split or split, get_dataset_split_names(source, subset) - ), - num_proc=None if args is None else args.num_proc, - ) - return d - - -def load_file(path: Path) -> Dataset: - logger.debug(f"Loading dataset from file: {path}") - if path.suffix not in _FORMATS: - raise ValueError(f"Unsupported file format: {path.suffix!r}") - - return load_dataset(_FORMATS[path.suffix], data_files=[str(path)], split="train") - - -def adaptive_load_dataset( - source: str, - split: str | None = None, - args: ProcArgs | None = None, -) -> Dataset: - path = Path(source) - - d = None - - # non-local -> HF remote - if not path.exists(): - if split is None: - raise ValueError("Split must be specified for remote datasets.") - d = load_remote(source, split, args) - - # local -> 1. dir 2. file - elif path.is_dir(): - d = load_local(source, split) - # try to load from locally hf-downloaded datasets - if d is None and split is not None: - d = load_remote(source, split, args) - - elif path.is_file(): - d = load_file(path) - - if d is None: - raise ValueError(f"Failed to load dataset: source={source!r}, split={split!r}") - - # not randomly shuffled - if args is not None and args.seed is not None: - d = d.shuffle(seed=None if args.seed < 0 else args.seed) - if args is not None and args.max_samples is not None: - d = d.select(range(min(args.max_samples, len(d)))) - return d + if isinstance(entry, PILImage.Image): + ext = (entry.format or "PNG").lower() + if ext == "jpeg": + ext = "jpg" + img = entry + if ext == "jpg" and img.mode not in ("RGB", "L"): + img = img.convert("RGB") + out = dest.with_suffix(f".{ext}") + img.save(out) + return out + + if isinstance(entry, dict): + raw, src_path = entry.get("bytes"), entry.get("path") + if raw: + out = dest.with_suffix(f".{guess_ext(raw, src_path)}") + out.write_bytes(raw) + return out + if src_path and Path(src_path).exists(): + src = Path(src_path) + out = dest.with_suffix(src.suffix.lower()) + shutil.copyfile(src, out) + return out + + if isinstance(entry, (str, Path)) and Path(entry).exists(): + src = Path(entry) + out = dest.with_suffix(src.suffix.lower()) + shutil.copyfile(src, out) + return out + except Exception as exc: + logger.warning(f"{WARN_PREFIX}Failed to dump preview image: {exc}") + return None @dataclass(frozen=True) @@ -328,12 +233,51 @@ class OutputActions(PathIO): hf_repo: str hf_subset: str | None hf_private: bool + + preview: int | None = None interactive: bool = False + def do_dump( + self, + d: "Dataset", + pipeline: FormatterPipeline, + id_override: str | None = None, + ) -> None: + """Dump images from the first ``self.preview`` rows for visual inspection. + + Since image columns are stored as opaque binary blobs in parquet/arrow + outputs, this writes a few representative image files to disk under the + default save path so an agent or user can directly inspect them. + + Args: + d: The formatted dataset to preview. + pipeline: The pipeline used to produce ``d`` (used for the output path). + id_override: Overrides the pipeline ID in the output path, mirroring + the behavior of ``OutputActions.do_save``. + """ + if self.preview is None or self.preview <= 0: + return + dump_dir = self.default_save_path( + pipeline, as_parquet=False, id_override=id_override + ) + dump_dir.mkdir(parents=True, exist_ok=True) + + for i in range(min(self.preview, len(d))): + try: + row = d[i] + except Exception as exc: + logger.warning( + f"{WARN_PREFIX}Failed to load sample {i} for preview: {exc}" + ) + continue + for j, img in enumerate(iter_images(row)): + dest = dump_dir / f"{i:06d}_{j:02d}" + path = _write_preview_image(img, dest) + if path is not None: + typer.echo(f"{PREVIEW_PREFIX}Preview image saved: {path}") + @staticmethod def _resolve_save_path(save_path: Path) -> Path | None: - import typer - def save_target(e: str) -> Path | None: if e.strip() == SKIP: return None @@ -369,8 +313,6 @@ def do_save( id_override: str | None = None, dry_run: bool = False, ): - import typer - save_path = self.default_save_path( pipeline, self.save_parquet, id_override=id_override ) diff --git a/src/prep/api/formatutils/__init__.py b/src/prep/api/formatutils/__init__.py index 11bbc92..779e2fa 100644 --- a/src/prep/api/formatutils/__init__.py +++ b/src/prep/api/formatutils/__init__.py @@ -1,4 +1,4 @@ -from .images import extract_images +from .images import extract_images, iter_images from .labels import extract_label from .options import ( extract_options, @@ -16,4 +16,5 @@ "format_options", "get_options_from_multi_entry", "get_options_from_single_entry", + "iter_images", ] diff --git a/src/prep/api/formatutils/images.py b/src/prep/api/formatutils/images.py index 5f9988f..b7d0cc0 100644 --- a/src/prep/api/formatutils/images.py +++ b/src/prep/api/formatutils/images.py @@ -1,3 +1,48 @@ +from collections.abc import Iterator +from typing import Any + +IMAGE_FIELDS = ("images", "image", "img") + + +def iter_images(e: dict) -> Iterator[Any]: + """Iterate over image entries in an example dictionary. + + Supports the same conventions as :func:`extract_images`: ``images``, + ``image``, ``img``, ``image_1``..``image_n``, and ``image_01``..``image_99``. + Values may be PIL images, raw ``{"bytes", "path"}`` records, or file paths, + depending on whether the dataset decodes its image feature. + + Args: + e: Example dictionary potentially containing image fields. + + Yields: + Image entries found in the example. + """ + if "image_1" in e: + for i in range(1, 100): + if f"image_{i}" not in e: + break + yield e[f"image_{i}"] + return + + if "image_01" in e: + for i in range(1, 100): + if f"image_{i:02d}" not in e: + break + yield e[f"image_{i:02d}"] + return + + for col in IMAGE_FIELDS: + value = e.get(col) + if value is None: + continue + if isinstance(value, (list, tuple)): + yield from value + else: + yield value + return + + def extract_images(e: dict) -> list: """Extract image paths from an example dictionary. diff --git a/src/prep/api/load.py b/src/prep/api/load.py new file mode 100644 index 0000000..3538c2d --- /dev/null +++ b/src/prep/api/load.py @@ -0,0 +1,156 @@ +import re +from pathlib import Path +from typing import TYPE_CHECKING + +from datasets import ( + Dataset, + get_dataset_split_names, + load_dataset, + load_from_disk, +) + +from .log import get_logger + +if TYPE_CHECKING: + from .types import ProcArgs + +logger = get_logger(__name__) + +_FORMATS = { + ".parquet": "parquet", + ".pq": "parquet", + ".json": "json", + ".jsonl": "json", + ".ndjson": "json", + ".csv": "csv", + ".tsv": "csv", + ".arrow": "arrow", + ".txt": "text", + ".text": "text", +} + + +def resolve_remote(source: str) -> tuple[str, str | None, str | None]: + """Resolve a remote dataset source into its components. + + Args: + source (str): dataset ID in the format of `org/data[@subset][:split]`. + + Returns: + tuple[str, str | None, str | None]: A tuple containing the source, subset, and split. + """ + match = re.match( + r"^(?P[^@]+)(?:@(?P[^:]+))?(?::(?P.+))?$", source + ) + if not match: + raise ValueError(f"Invalid source format: {source!r}") + + return ( + match.group("source"), + match.group("subset"), + match.group("split"), + ) + + +def resolve_split(split: str | None, available: list[str]) -> str: + if split is None: + if len(available) == 1: + return available[0] + raise ValueError( + f"Split must be specified for dataset with multiple splits: {available}" + ) + + # direct hit is preferred + if split in available: + return split + + # translate if "val" -> "validation" if exists + # we want to be generous here to facilitate IO + if split == "val" and "validation" in available: + return "validation" + + raise ValueError(f"Split {split!r} not found in dataset: {available}") + + +def load_local(source: str, split: str | None = None) -> Dataset | None: + logger.debug(f"Loading dataset from local disk: {source} (split={split})") + try: + d = load_from_disk(source) + except Exception: + return None + + if isinstance(d, Dataset): + return d + + return d[resolve_split(split, list(map(str, d.keys())))] + + +def load_remote( + source: str, split: str | None = None, args: "ProcArgs | None" = None +) -> Dataset: + """Load a dataset from Hugging Face Hub. + + Args: + source (str): dataset ID in the format of `org/data[@subset][:split]`. + split (str | None, optional): dataset split name. Defaults to None. + args (ProcArgs | None, optional): processing arguments. Defaults to None. + + Returns: + Dataset: the loaded dataset. + """ + logger.debug(f"Loading dataset from Hugging Face Hub: {source} (split={split})") + source, subset, override_split = resolve_remote(source) + + d = load_dataset( + source, + subset, + split=resolve_split( + override_split or split, get_dataset_split_names(source, subset) + ), + num_proc=None if args is None else args.num_proc, + ) + return d + + +def load_file(path: Path) -> Dataset: + logger.debug(f"Loading dataset from file: {path}") + if path.suffix not in _FORMATS: + raise ValueError(f"Unsupported file format: {path.suffix!r}") + + return load_dataset(_FORMATS[path.suffix], data_files=[str(path)], split="train") + + +def adaptive_load_dataset( + source: str, + split: str | None = None, + args: "ProcArgs | None" = None, +) -> Dataset: + path = Path(source) + + d = None + + # non-local -> HF remote + if not path.exists(): + if split is None: + raise ValueError("Split must be specified for remote datasets.") + d = load_remote(source, split, args) + + # local -> 1. dir 2. file + elif path.is_dir(): + d = load_local(source, split) + # try to load from locally hf-downloaded datasets + if d is None and split is not None: + d = load_remote(source, split, args) + + elif path.is_file(): + d = load_file(path) + + if d is None: + raise ValueError(f"Failed to load dataset: source={source!r}, split={split!r}") + + # not randomly shuffled + if args is not None and args.seed is not None: + d = d.shuffle(seed=None if args.seed < 0 else args.seed) + if args is not None and args.max_samples is not None: + d = d.select(range(min(args.max_samples, len(d)))) + return d diff --git a/src/prep/cli.py b/src/prep/cli.py index 3dc35ba..32432a8 100644 --- a/src/prep/cli.py +++ b/src/prep/cli.py @@ -50,6 +50,12 @@ def prep( save_nproc: int | None = typer.Option( None, envvar="SAVE_NPROC", rich_help_panel=_SAVE ), + save_preview: int = typer.Option( + 0, + envvar="SAVE_PREVIEW", + rich_help_panel=_SAVE, + help="Save images from the first N samples for inspection.", + ), hf: bool | None = typer.Option(None, envvar="HF", rich_help_panel=_HF), hf_repo: str | None = typer.Option(None, envvar="HF_REPO", rich_help_panel=_HF), hf_subset: str | None = typer.Option(None, envvar="HF_SUBSET", rich_help_panel=_HF), @@ -125,8 +131,10 @@ def prep( hf_repo=hf_repo or pipeline_id, hf_subset=hf_subset, hf_private=hf_private, + preview=save_preview, interactive=interactive, ) + action.do_dump(d, pipeline=pipeline, id_override=pipeline_id) action.do_save(d, pipeline=pipeline, nproc=save_nproc, id_override=pipeline_id) action.do_upload(d, split=pipeline.split, nproc=hf_nproc) diff --git a/src/prep/constants.py b/src/prep/constants.py index c5edd3f..332d34b 100644 --- a/src/prep/constants.py +++ b/src/prep/constants.py @@ -25,6 +25,7 @@ def is_env_enabled(env_var: str, default: str = "0") -> bool: SAVE_PREFIX = "💾\t" HF_PREFIX = "☁️\t" +PREVIEW_PREFIX = "🖼️\t" SKIP = "SKIP" WARN_PREFIX = "⚠️\t" diff --git a/tests/api/formatutils/test_images.py b/tests/api/formatutils/test_images.py index 538a729..3a0d5f3 100644 --- a/tests/api/formatutils/test_images.py +++ b/tests/api/formatutils/test_images.py @@ -1,6 +1,6 @@ import pytest -from prep.api.formatutils import extract_images +from prep.api.formatutils import extract_images, iter_images @pytest.mark.parametrize( @@ -72,3 +72,29 @@ def test_extract_images_handles_priority_and_sequence_boundaries( expected: list, ) -> None: assert extract_images(entry) == expected + + +@pytest.mark.parametrize( + ("entry", "expected"), + [ + pytest.param({"image": "a.png"}, ["a.png"], id="single-image"), + pytest.param({"images": ["a.png", "b.png"]}, ["a.png", "b.png"], id="images"), + pytest.param({"img": "a.png"}, ["a.png"], id="img"), + pytest.param( + {"image_1": "a.png", "image_2": "b.png", "image_4": "skip.png"}, + ["a.png", "b.png"], + id="numbered-stops-at-gap", + ), + pytest.param( + {"image_01": "a.png", "image_02": "b.png"}, + ["a.png", "b.png"], + id="zero-padded", + ), + pytest.param({"image": None}, [], id="none-value-yields-nothing"), + pytest.param({"question": "q"}, [], id="no-image-field"), + ], +) +def test_iter_images_yields_entries_across_conventions( + entry: dict, expected: list +) -> None: + assert list(iter_images(entry)) == expected diff --git a/tests/api/test_dataio.py b/tests/api/test_dataio.py index defe42d..05a2328 100644 --- a/tests/api/test_dataio.py +++ b/tests/api/test_dataio.py @@ -1,19 +1,21 @@ +from io import BytesIO from pathlib import Path import pytest -from datasets import Dataset, DatasetDict +from datasets import Dataset, Features, Image, List +from PIL import Image as PILImage from prep.api.dataio import ( OutputActions, PathIO, - adaptive_load_dataset, - load_local, - load_remote, - resolve_remote, - resolve_split, ) from prep.api.formatter import FormatterPipeline -from prep.api.types import ProcArgs + + +def _png_bytes(size: int = 8) -> bytes: + buf = BytesIO() + PILImage.new("RGB", (size, size), "red").save(buf, format="PNG") + return buf.getvalue() def test_status_dicts_reports_registered_and_unregistered_datasets( @@ -66,133 +68,6 @@ def test_status_dicts_reports_registered_and_unregistered_datasets( assert unregistered_report["test"] == [True] -class TestResolveRemote: - """Tests for the resolve_remote() function.""" - - @pytest.mark.parametrize( - ("remote", "expected"), - [ - ("org/dataset", ("org/dataset", None, None)), - ("org/dataset@train", ("org/dataset", "train", None)), - ("org/dataset:validation", ("org/dataset:validation", None, None)), - ("org/dataset@config:test", ("org/dataset", "config", "test")), - ("org/team/dataset", ("org/team/dataset", None, None)), - ("org/team/dataset@subset1:val", ("org/team/dataset", "subset1", "val")), - ( - "org/dataset@config:train-splits-001", - ("org/dataset", "config", "train-splits-001"), - ), - (":", (":", None, None)), - ("org/@dataset@extra:test", ("org/", "dataset@extra", "test")), - ("org/dataset@sub:part1:part2", ("org/dataset", "sub", "part1:part2")), - ], - ) - def test_resolve_remote( - self, - remote: str, - expected: tuple[str, str | None, str | None], - ) -> None: - """Test valid remote source parsing.""" - assert resolve_remote(remote) == expected - - @pytest.mark.parametrize( - "remote", - [ - pytest.param("", id="empty-string"), - pytest.param("@", id="at-sign-alone"), - ], - ) - def test_invalid_source_raises(self, remote: str) -> None: - """Test invalid source formats raise ValueError.""" - with pytest.raises(ValueError, match="Invalid source format"): - resolve_remote(remote) - - -class TestResolveSplit: - @pytest.mark.parametrize( - ("requested", "available", "expected"), - [ - (None, ["train"], "train"), - ("test", ["train", "test"], "test"), - ("val", ["validation"], "validation"), - ], - ) - def test_accepts_single_direct_and_validation_alias( - self, requested: str | None, available: list[str], expected: str - ) -> None: - assert resolve_split(requested, available) == expected - - def test_rejects_ambiguous_or_missing_splits(self) -> None: - with pytest.raises(ValueError, match="must be specified"): - resolve_split(None, ["train", "test"]) - with pytest.raises(ValueError, match="not found"): - resolve_split("dev", ["train"]) - - -def test_load_local_handles_dataset_dict_and_unloadable_path( - monkeypatch: pytest.MonkeyPatch, -) -> None: - dataset = Dataset.from_dict({"value": [1]}) - monkeypatch.setattr( - "prep.api.dataio.load_from_disk", lambda _: DatasetDict({"train": dataset}) - ) - assert load_local("ignored", "train") == dataset - - monkeypatch.setattr( - "prep.api.dataio.load_from_disk", - lambda _: (_ for _ in ()).throw(OSError("bad")), - ) - assert load_local("ignored") is None - - -def test_load_remote_parses_source_and_passes_processing_options( - monkeypatch: pytest.MonkeyPatch, -) -> None: - expected = Dataset.from_dict({"value": [1]}) - calls: dict[str, object] = {} - monkeypatch.setattr( - "prep.api.dataio.get_dataset_split_names", lambda *_: ["validation"] - ) - monkeypatch.setattr( - "prep.api.dataio.load_dataset", - lambda *args, **kwargs: calls.update(args=args, **kwargs) or expected, - ) - - assert load_remote("org/data@subset:val", args=ProcArgs(num_proc=2)) == expected - assert calls == { - "args": ("org/data", "subset"), - "split": "validation", - "num_proc": 2, - } - - -def test_adaptive_load_dataset_loads_files_and_applies_sampling( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - source = tmp_path / "data.csv" - source.write_text("value\n1\n2\n3\n", encoding="utf-8") - dataset = Dataset.from_dict({"value": [1, 2, 3]}) - monkeypatch.setattr("prep.api.dataio.load_file", lambda _: dataset) - - result = adaptive_load_dataset(str(source), args=ProcArgs(seed=3, max_samples=2)) - assert len(result) == 2 - assert set(result["value"]).issubset({1, 2, 3}) - - -def test_adaptive_load_dataset_requires_remote_split_and_falls_back_to_remote( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - with pytest.raises(ValueError, match="Split must be specified"): - adaptive_load_dataset("org/missing") - - local_dir = tmp_path / "downloaded" - local_dir.mkdir() - expected = Dataset.from_dict({"value": [1]}) - monkeypatch.setattr("prep.api.dataio.load_local", lambda *_: None) - monkeypatch.setattr("prep.api.dataio.load_remote", lambda *_: expected) - assert adaptive_load_dataset(str(local_dir), split="train") == expected - - def test_pathio_helpers_and_tables(tmp_path: Path) -> None: dataset_dir = tmp_path / "sft" / "demo" assert ( @@ -254,3 +129,108 @@ def test_output_actions_skip_non_interactive_operations( action.do_save(dataset, pipeline) action.do_upload(dataset, "train") assert "Skipping (non-interactive mode)" in capsys.readouterr().out + + +class TestDoDump: + """Tests for OutputActions.do_dump, which writes preview images to disk.""" + + @staticmethod + def _action(save_root: Path, preview: int) -> OutputActions: + return OutputActions( + save_root=save_root, + save=False, + save_parquet=False, + hf=False, + hf_repo="org/repo", + hf_subset=None, + hf_private=True, + preview=preview, + ) + + @staticmethod + def _blob_image_dataset(sizes: list[int]) -> Dataset: + return Dataset.from_dict( + {"image": [_png_bytes(size) for size in sizes]}, + features=Features(image=Image(decode=False)), + ) + + def test_dumps_undecoded_blob_images( + self, tmp_path: Path, capsys: pytest.CaptureFixture + ) -> None: + dataset = self._blob_image_dataset([9, 10]) + pipeline = FormatterPipeline("demo", "cls", "train", lambda *_: dataset) + + self._action(tmp_path, preview=2).do_dump(dataset, pipeline) + + dump_dir = tmp_path / "cls" / "demo" / "train" + assert sorted(p.name for p in dump_dir.iterdir()) == [ + "000000_00.png", + "000001_00.png", + ] + for path in dump_dir.iterdir(): + with PILImage.open(path) as img: + img.verify() + assert capsys.readouterr().out.count("Preview image saved") == 2 + + def test_dumps_decoded_pil_and_multi_image_rows(self, tmp_path: Path) -> None: + dataset = Dataset.from_dict( + {"images": [[_png_bytes(8)], [_png_bytes(9), _png_bytes(10)]]}, + features=Features(images=List(Image(decode=True))), + ) + pipeline = FormatterPipeline("demo", "sft", "train", lambda *_: dataset) + + self._action(tmp_path, preview=2).do_dump(dataset, pipeline) + + dump_dir = tmp_path / "sft" / "demo" / "train" + assert sorted(p.name for p in dump_dir.iterdir()) == [ + "000000_00.png", + "000001_00.png", + "000001_01.png", + ] + with PILImage.open(dump_dir / "000001_01.png") as img: + assert img.size == (10, 10) + + def test_no_op_for_nonpositive_preview(self, tmp_path: Path) -> None: + dataset = self._blob_image_dataset([8]) + pipeline = FormatterPipeline("demo", "cls", "train", lambda *_: dataset) + + self._action(tmp_path, preview=0).do_dump(dataset, pipeline) + + assert list(tmp_path.iterdir()) == [] + + def test_no_op_for_empty_dataset(self, tmp_path: Path) -> None: + dataset = Dataset.from_dict({"image": []}) + pipeline = FormatterPipeline("demo", "cls", "train", lambda *_: dataset) + + self._action(tmp_path, preview=3).do_dump(dataset, pipeline) + + assert not [p for p in tmp_path.rglob("*") if p.is_file()] + + def test_clamps_preview_count_and_honors_id_override(self, tmp_path: Path) -> None: + dataset = self._blob_image_dataset([8]) + pipeline = FormatterPipeline("demo", "cls", "train", lambda *_: dataset) + + self._action(tmp_path, preview=10).do_dump( + dataset, pipeline, id_override="alias" + ) + + assert [p.name for p in (tmp_path / "cls" / "alias" / "train").iterdir()] == [ + "000000_00.png" + ] + + def test_skips_unwritable_entries( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture, + ) -> None: + dataset = self._blob_image_dataset([8]) + pipeline = FormatterPipeline("demo", "cls", "train", lambda *_: dataset) + monkeypatch.setattr( + "prep.api.dataio._write_preview_image", lambda entry, dest: None + ) + + self._action(tmp_path, preview=1).do_dump(dataset, pipeline) + + assert not [p for p in tmp_path.rglob("*") if p.is_file()] + assert "Preview image saved" not in capsys.readouterr().out diff --git a/tests/api/test_load.py b/tests/api/test_load.py new file mode 100644 index 0000000..5f46dd1 --- /dev/null +++ b/tests/api/test_load.py @@ -0,0 +1,151 @@ +from pathlib import Path + +import pytest +from datasets import Dataset, DatasetDict + +from prep.api.load import ( + adaptive_load_dataset, + load_local, + load_remote, + resolve_remote, + resolve_split, +) +from prep.api.types import ProcArgs + + +class TestResolveRemote: + """resolve_remote() parses ``org/data[@subset][:split]`` sources.""" + + @pytest.mark.parametrize( + ("remote", "expected"), + [ + ("org/dataset", ("org/dataset", None, None)), + ("org/dataset@train", ("org/dataset", "train", None)), + ("org/dataset:validation", ("org/dataset:validation", None, None)), + ("org/dataset@config:test", ("org/dataset", "config", "test")), + ("org/team/dataset", ("org/team/dataset", None, None)), + ("org/team/dataset@subset1:val", ("org/team/dataset", "subset1", "val")), + ( + "org/dataset@config:train-splits-001", + ("org/dataset", "config", "train-splits-001"), + ), + (":", (":", None, None)), + ("org/@dataset@extra:test", ("org/", "dataset@extra", "test")), + ("org/dataset@sub:part1:part2", ("org/dataset", "sub", "part1:part2")), + ], + ) + def test_resolve_remote( + self, + remote: str, + expected: tuple[str, str | None, str | None], + ) -> None: + assert resolve_remote(remote) == expected + + @pytest.mark.parametrize("remote", ["", "@"]) + def test_invalid_source_raises(self, remote: str) -> None: + with pytest.raises(ValueError, match="Invalid source format"): + resolve_remote(remote) + + +class TestResolveSplit: + """resolve_split() matches requested splits against available ones.""" + + @pytest.mark.parametrize( + ("requested", "available", "expected"), + [ + (None, ["train"], "train"), + ("test", ["train", "test"], "test"), + ("val", ["validation"], "validation"), + ], + ) + def test_resolves_requested_split( + self, requested: str | None, available: list[str], expected: str + ) -> None: + assert resolve_split(requested, available) == expected + + def test_ambiguous_default_requires_explicit_split(self) -> None: + with pytest.raises(ValueError, match="must be specified"): + resolve_split(None, ["train", "test"]) + + def test_unknown_split_raises(self) -> None: + with pytest.raises(ValueError, match="not found"): + resolve_split("dev", ["train"]) + + +class TestLoadLocal: + """load_local() wraps datasets.load_from_disk with split resolution.""" + + def test_selects_split_from_dataset_dict( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + dataset = Dataset.from_dict({"value": [1]}) + monkeypatch.setattr( + "prep.api.load.load_from_disk", + lambda _: DatasetDict({"train": dataset}), + ) + assert load_local("ignored", "train") == dataset + + def test_returns_none_when_path_is_unloadable( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def raise_oserror(_: object) -> Dataset: + raise OSError("bad") + + monkeypatch.setattr("prep.api.load.load_from_disk", raise_oserror) + assert load_local("ignored") is None + + +class TestLoadRemote: + """load_remote() resolves sources and delegates to datasets.load_dataset.""" + + def test_parses_source_and_passes_processing_options( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + expected = Dataset.from_dict({"value": [1]}) + calls: dict[str, object] = {} + monkeypatch.setattr( + "prep.api.load.get_dataset_split_names", lambda *_: ["validation"] + ) + monkeypatch.setattr( + "prep.api.load.load_dataset", + lambda *args, **kwargs: calls.update(args=args, **kwargs) or expected, + ) + + assert load_remote("org/data@subset:val", args=ProcArgs(num_proc=2)) == expected + assert calls == { + "args": ("org/data", "subset"), + "split": "validation", + "num_proc": 2, + } + + +class TestAdaptiveLoadDataset: + """adaptive_load_dataset() picks local/remote loaders and applies sampling.""" + + def test_loads_local_file_and_applies_sampling( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + source = tmp_path / "data.csv" + source.write_text("value\n1\n2\n3\n", encoding="utf-8") + dataset = Dataset.from_dict({"value": [1, 2, 3]}) + monkeypatch.setattr("prep.api.load.load_file", lambda _: dataset) + + result = adaptive_load_dataset( + str(source), args=ProcArgs(seed=3, max_samples=2) + ) + assert len(result) == 2 + assert set(result["value"]).issubset({1, 2, 3}) + + def test_remote_source_requires_split(self) -> None: + with pytest.raises(ValueError, match="Split must be specified"): + adaptive_load_dataset("org/missing") + + def test_local_dir_falls_back_to_remote( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + local_dir = tmp_path / "downloaded" + local_dir.mkdir() + expected = Dataset.from_dict({"value": [1]}) + monkeypatch.setattr("prep.api.load.load_local", lambda *_: None) + monkeypatch.setattr("prep.api.load.load_remote", lambda *_: expected) + assert adaptive_load_dataset(str(local_dir), split="train") == expected From 574cfb1d96595aeea69d378de01dfdb3cb84fbd9 Mon Sep 17 00:00:00 2001 From: Moenupa <49304833+Moenupa@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:11:44 +0800 Subject: [PATCH 2/2] [fix] implment load_remote org/data[@subset][:split] --- src/prep/api/load.py | 6 ++++-- tests/api/test_load.py | 21 ++++++++++++++++++--- tests/transform/test_crop.py | 27 ++++++++++++++++++++++++++- 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/prep/api/load.py b/src/prep/api/load.py index 3538c2d..85e2261 100644 --- a/src/prep/api/load.py +++ b/src/prep/api/load.py @@ -40,7 +40,8 @@ def resolve_remote(source: str) -> tuple[str, str | None, str | None]: tuple[str, str | None, str | None]: A tuple containing the source, subset, and split. """ match = re.match( - r"^(?P[^@]+)(?:@(?P[^:]+))?(?::(?P.+))?$", source + r"^(?P[a-zA-Z\d/._-]+)(?:@(?P[^:]+))?(?::(?P.+))?$", + source, ) if not match: raise ValueError(f"Invalid source format: {source!r}") @@ -131,7 +132,8 @@ def adaptive_load_dataset( # non-local -> HF remote if not path.exists(): - if split is None: + _, _, inline_split = resolve_remote(source) + if split is None and inline_split is None: raise ValueError("Split must be specified for remote datasets.") d = load_remote(source, split, args) diff --git a/tests/api/test_load.py b/tests/api/test_load.py index 5f46dd1..54b7b89 100644 --- a/tests/api/test_load.py +++ b/tests/api/test_load.py @@ -21,7 +21,7 @@ class TestResolveRemote: [ ("org/dataset", ("org/dataset", None, None)), ("org/dataset@train", ("org/dataset", "train", None)), - ("org/dataset:validation", ("org/dataset:validation", None, None)), + ("org/dataset:validation", ("org/dataset", None, "validation")), ("org/dataset@config:test", ("org/dataset", "config", "test")), ("org/team/dataset", ("org/team/dataset", None, None)), ("org/team/dataset@subset1:val", ("org/team/dataset", "subset1", "val")), @@ -29,7 +29,6 @@ class TestResolveRemote: "org/dataset@config:train-splits-001", ("org/dataset", "config", "train-splits-001"), ), - (":", (":", None, None)), ("org/@dataset@extra:test", ("org/", "dataset@extra", "test")), ("org/dataset@sub:part1:part2", ("org/dataset", "sub", "part1:part2")), ], @@ -41,7 +40,7 @@ def test_resolve_remote( ) -> None: assert resolve_remote(remote) == expected - @pytest.mark.parametrize("remote", ["", "@"]) + @pytest.mark.parametrize("remote", ["", "@", ":"]) def test_invalid_source_raises(self, remote: str) -> None: with pytest.raises(ValueError, match="Invalid source format"): resolve_remote(remote) @@ -140,6 +139,22 @@ def test_remote_source_requires_split(self) -> None: with pytest.raises(ValueError, match="Split must be specified"): adaptive_load_dataset("org/missing") + def test_remote_source_with_inline_split( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + expected = Dataset.from_dict({"value": [1]}) + calls: list[tuple[str, str | None]] = [] + + def fake_load_remote( + source: str, split: str | None = None, args: object = None + ) -> Dataset: + calls.append((source, split)) + return expected + + monkeypatch.setattr("prep.api.load.load_remote", fake_load_remote) + assert adaptive_load_dataset("org/data@subset:train") == expected + assert calls == [("org/data@subset:train", None)] + def test_local_dir_falls_back_to_remote( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/transform/test_crop.py b/tests/transform/test_crop.py index 2fd8119..2821a18 100644 --- a/tests/transform/test_crop.py +++ b/tests/transform/test_crop.py @@ -1,6 +1,9 @@ +from pathlib import Path + +import pytest from PIL import Image -from prep.transform.crop import crop_black_border +from prep.transform.crop import crop_black_border, crop_black_columns def _framed( @@ -15,6 +18,11 @@ def _framed( return img +@pytest.fixture +def real_image_path() -> Path: + return Path("in/test.jpg") + + class TestCropBlackBorder: def test_crops_exact_black_frame(self) -> None: assert crop_black_border(_framed()).size == (4, 4) @@ -54,3 +62,20 @@ def test_preserves_content_and_mode(self) -> None: rgba = _framed().convert("RGBA") assert crop_black_border(rgba).mode == "RGBA" + + def test_real_case(self, real_image_path: Path) -> None: + try: + img = Image.open(real_image_path) + except Exception as e: + pytest.skip(f"Real image case not available {real_image_path}: {e}") + + crop_black_border(img).save( + real_image_path.with_name( + real_image_path.stem + "_cropped" + real_image_path.suffix + ) + ) + crop_black_columns(img).save( + real_image_path.with_name( + real_image_path.stem + "_cropped_columns" + real_image_path.suffix + ) + )