Frank Schwenk
a60a18a253
Introduce ruff lint/format config, expand .gitignore, and reformat Python sources. Co-authored-by: Cursor <cursoragent@cursor.com>
242 lines
8.4 KiB
Python
242 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from imagepipeline.core.pipeline import Pipeline
|
|
from imagepipeline.core.resume import (
|
|
expected_output_filenames,
|
|
materialize_external_outputs,
|
|
step_outputs_complete,
|
|
)
|
|
from imagepipeline.core.step import StepDefinition
|
|
from imagepipeline.modules.imagemagick_grayscale import ImageMagickGrayscale
|
|
from imagepipeline.modules.rembg import RembgModule
|
|
from imagepipeline.utils.gmic import finalize_gmic_output, split_gmic_command
|
|
from tests.conftest import make_png
|
|
|
|
|
|
class TestGmicCommandSplit:
|
|
def test_splits_command_and_arguments(self) -> None:
|
|
assert split_gmic_command("-gcd_stereo_img 0,0,2.028,1,1.714,3.06,4,1,0") == [
|
|
"-gcd_stereo_img",
|
|
"0,0,2.028,1,1.714,3.06,4,1,0",
|
|
]
|
|
|
|
def test_preserves_quoted_empty_argument(self) -> None:
|
|
parts = split_gmic_command('-fx_custom_gradient 0,0,0,"",1,0')
|
|
assert parts == ["-fx_custom_gradient", "0,0,0,,1,0"]
|
|
|
|
|
|
class TestFinalizeGmicOutput:
|
|
def test_keeps_frame_000001_and_removes_000000(self, tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "out"
|
|
output_dir.mkdir()
|
|
intended = output_dir / "photo.png"
|
|
frame_000000 = output_dir / "photo_000000.png"
|
|
frame_000001 = output_dir / "photo_000001.png"
|
|
frame_000000.write_bytes(b"discard")
|
|
frame_000001.write_bytes(b"keep")
|
|
|
|
result = finalize_gmic_output(output_dir, intended)
|
|
|
|
assert result == intended
|
|
assert intended.read_bytes() == b"keep"
|
|
assert not frame_000000.exists()
|
|
assert not frame_000001.exists()
|
|
|
|
def test_leaves_single_output_unchanged(self, tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "out"
|
|
output_dir.mkdir()
|
|
intended = output_dir / "photo.png"
|
|
intended.write_bytes(b"single")
|
|
|
|
result = finalize_gmic_output(output_dir, intended)
|
|
|
|
assert result == intended
|
|
assert intended.read_bytes() == b"single"
|
|
|
|
def test_renames_only_000000_when_000001_missing(self, tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "out"
|
|
output_dir.mkdir()
|
|
intended = output_dir / "photo.png"
|
|
frame_000000 = output_dir / "photo_000000.png"
|
|
frame_000000.write_bytes(b"only")
|
|
|
|
result = finalize_gmic_output(output_dir, intended)
|
|
|
|
assert result == intended
|
|
assert intended.read_bytes() == b"only"
|
|
assert not frame_000000.exists()
|
|
|
|
|
|
class TestExpectedOutputFilenames:
|
|
def test_rembg_maps_jpg_inputs_to_png_outputs(self, tmp_path: Path) -> None:
|
|
jpg = tmp_path / "photo.jpg"
|
|
jpg.write_bytes(b"jpeg")
|
|
step = StepDefinition(
|
|
step_id="rembg_01",
|
|
module_name="rembg",
|
|
module=RembgModule,
|
|
input_refs=["input"],
|
|
params={},
|
|
output_dir_name="rembg_01",
|
|
)
|
|
|
|
names = expected_output_filenames(
|
|
step,
|
|
matched_groups=[[jpg]],
|
|
input_paths=[jpg],
|
|
params=RembgModule.validate_module_params({}),
|
|
)
|
|
|
|
assert names == ["photo.png"]
|
|
|
|
def test_rembg_resume_detects_existing_png_outputs(self, tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "rembg_01"
|
|
output_dir.mkdir()
|
|
png = output_dir / "photo.png"
|
|
make_png(png)
|
|
jpg = tmp_path / "input" / "photo.jpg"
|
|
jpg.parent.mkdir()
|
|
jpg.write_bytes(b"jpeg")
|
|
step = StepDefinition(
|
|
step_id="rembg_01",
|
|
module_name="rembg",
|
|
module=RembgModule,
|
|
input_refs=["input"],
|
|
params={},
|
|
output_dir_name="rembg_01",
|
|
)
|
|
params = RembgModule.validate_module_params({})
|
|
expected = expected_output_filenames(
|
|
step,
|
|
matched_groups=[[jpg]],
|
|
input_paths=[jpg],
|
|
params=params,
|
|
)
|
|
|
|
assert step_outputs_complete([output_dir / name for name in expected])
|
|
|
|
|
|
class TestPipelineResume:
|
|
@pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed")
|
|
def test_continue_skips_completed_steps(
|
|
self, input_dir: Path, output_base: Path, capsys
|
|
) -> None:
|
|
with Pipeline(
|
|
name="resume_test",
|
|
input_dir=input_dir,
|
|
output_base=output_base,
|
|
verbose=True,
|
|
) as p:
|
|
first = p.step("imagemagick_grayscale", inputs="input")
|
|
p.step("imagemagick_grayscale", inputs=first)
|
|
root = p.run()
|
|
|
|
capsys.readouterr()
|
|
|
|
with Pipeline(
|
|
name="resume_test",
|
|
input_dir=input_dir,
|
|
output_base=output_base,
|
|
verbose=True,
|
|
continue_from=root,
|
|
) as p:
|
|
step_a = p.step("imagemagick_grayscale", inputs="input")
|
|
p.step("imagemagick_grayscale", inputs=step_a)
|
|
resumed_root = p.run()
|
|
|
|
assert resumed_root == root
|
|
output = capsys.readouterr().out
|
|
assert "Skipped step imagemagick_grayscale_01" in output
|
|
assert "Skipped step imagemagick_grayscale_02" in output
|
|
|
|
@pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed")
|
|
def test_existing_outputs_reuse_external_folder(
|
|
self, input_dir: Path, output_base: Path, tmp_path: Path, capsys
|
|
) -> None:
|
|
external = tmp_path / "external_rembg"
|
|
external.mkdir()
|
|
for src in input_dir.iterdir():
|
|
if src.is_file():
|
|
make_png(external / src.name, width=4, height=4, rgb=(10, 20, 30))
|
|
|
|
with Pipeline(
|
|
name="external_test",
|
|
input_dir=input_dir,
|
|
output_base=output_base,
|
|
verbose=True,
|
|
existing_outputs={"imagemagick_grayscale_01": external},
|
|
) as p:
|
|
reused = p.step("imagemagick_grayscale", inputs="input")
|
|
p.step("imagemagick_grayscale", inputs=reused)
|
|
root = p.run()
|
|
|
|
output = capsys.readouterr().out
|
|
assert "Reused external output for imagemagick_grayscale_01" in output
|
|
assert (root / "imagemagick_grayscale_01" / "photo_a.png").exists()
|
|
|
|
@pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed")
|
|
def test_existing_outputs_with_custom_step_id(
|
|
self, input_dir: Path, output_base: Path, tmp_path: Path, capsys
|
|
) -> None:
|
|
external = tmp_path / "external_gray"
|
|
external.mkdir()
|
|
for src in input_dir.iterdir():
|
|
if src.is_file():
|
|
make_png(external / src.name, width=4, height=4, rgb=(10, 20, 30))
|
|
|
|
with Pipeline(
|
|
name="custom_id_external",
|
|
input_dir=input_dir,
|
|
output_base=output_base,
|
|
verbose=True,
|
|
existing_outputs={"input_bokeh": external},
|
|
) as p:
|
|
reused = p.step("imagemagick_grayscale", inputs="input", step_id="input_bokeh")
|
|
p.step("imagemagick_grayscale", inputs=reused)
|
|
root = p.run()
|
|
|
|
output = capsys.readouterr().out
|
|
assert "Reused external output for input_bokeh" in output
|
|
assert (root / "input_bokeh" / "photo_a.png").exists()
|
|
assert not (root / "imagemagick_grayscale_01").exists()
|
|
|
|
|
|
class TestMaterializeExternal:
|
|
def test_links_files_by_stem(self, tmp_path: Path) -> None:
|
|
external = tmp_path / "external"
|
|
external.mkdir()
|
|
make_png(external / "photo_a.png")
|
|
output_dir = tmp_path / "out"
|
|
step = StepDefinition(
|
|
step_id="imagemagick_grayscale_01",
|
|
module_name="imagemagick_grayscale",
|
|
module=ImageMagickGrayscale,
|
|
input_refs=["input"],
|
|
params={},
|
|
output_dir_name="imagemagick_grayscale_01",
|
|
)
|
|
input_paths = [tmp_path / "photo_a.png"]
|
|
paths = materialize_external_outputs(
|
|
external,
|
|
output_dir,
|
|
step,
|
|
matched_groups=[[path] for path in input_paths],
|
|
input_paths=input_paths,
|
|
params={},
|
|
)
|
|
assert len(paths) == 1
|
|
assert paths[0].name == "photo_a.png"
|
|
assert paths[0].is_symlink()
|
|
|
|
def test_step_outputs_complete(self, tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "done"
|
|
output_dir.mkdir()
|
|
make_png(output_dir / "photo.png")
|
|
assert step_outputs_complete([output_dir / "photo.png"])
|
|
assert not step_outputs_complete([output_dir / "missing.png"])
|