You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
import pytest
|
|
from lightning_fabric import seed_everything
|
|
|
|
from imaginairy.img_processors.control_modes import CONTROL_MODES
|
|
from imaginairy.img_utils import pillow_img_to_torch_image, torch_img_to_pillow_img
|
|
from imaginairy.schema import LazyLoadingImage
|
|
from tests import TESTS_FOLDER
|
|
from tests.utils import assert_image_similar_to_expectation
|
|
|
|
|
|
def control_img_to_pillow_img(img_t):
|
|
return torch_img_to_pillow_img((img_t - 0.5) * 2)
|
|
|
|
|
|
control_mode_params = list(CONTROL_MODES.items())
|
|
|
|
|
|
@pytest.mark.parametrize(("control_name", "control_func"), control_mode_params)
|
|
def test_control_images(filename_base_for_outputs, control_func, control_name):
|
|
seed_everything(42)
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png")
|
|
img_t = pillow_img_to_torch_image(img)
|
|
if control_name == "inpaint":
|
|
control_t = control_func(img_t.clone(), img_t.clone())
|
|
else:
|
|
control_t = control_func(img_t.clone())
|
|
control_img = control_img_to_pillow_img(control_t)
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
|
|
assert_image_similar_to_expectation(control_img, img_path, threshold=8000)
|