IOPaint/lama_cleaner/tests/test_model.py

231 lines
6.8 KiB
Python
Raw Normal View History

2022-09-15 16:21:27 +02:00
import os
2022-04-15 18:11:51 +02:00
from pathlib import Path
import cv2
import pytest
2022-09-05 07:07:25 +02:00
import torch
2022-04-15 18:11:51 +02:00
from lama_cleaner.model_manager import ModelManager
2022-09-22 06:38:32 +02:00
from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler
2022-04-15 18:11:51 +02:00
current_dir = Path(__file__).parent.absolute().resolve()
2022-09-30 15:39:23 +02:00
save_dir = current_dir / 'result'
save_dir.mkdir(exist_ok=True, parents=True)
2022-09-05 07:07:25 +02:00
device = 'cuda' if torch.cuda.is_available() else 'cpu'
2022-04-15 18:11:51 +02:00
2022-09-15 16:21:27 +02:00
def get_data(fx=1, fy=1.0, img_p=current_dir / "image.png", mask_p=current_dir / "mask.png"):
img = cv2.imread(str(img_p))
2022-04-15 18:11:51 +02:00
img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
2022-09-15 16:21:27 +02:00
mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
2022-08-22 17:24:02 +02:00
2022-07-21 16:09:10 +02:00
if fx != 1:
2022-09-02 04:37:30 +02:00
img = cv2.resize(img, None, fx=fx, fy=fy, interpolation=cv2.INTER_AREA)
mask = cv2.resize(mask, None, fx=fx, fy=fy, interpolation=cv2.INTER_NEAREST)
2022-04-15 18:11:51 +02:00
return img, mask
2022-07-14 10:49:03 +02:00
def get_config(strategy, **kwargs):
data = dict(
2022-04-15 18:11:51 +02:00
ldm_steps=1,
2022-07-14 10:49:03 +02:00
ldm_sampler=LDMSampler.plms,
2022-04-15 18:11:51 +02:00
hd_strategy=strategy,
hd_strategy_crop_margin=32,
hd_strategy_crop_trigger_size=200,
hd_strategy_resize_limit=200,
)
2022-07-14 10:49:03 +02:00
data.update(**kwargs)
return Config(**data)
2022-04-15 18:11:51 +02:00
2022-09-15 16:21:27 +02:00
def assert_equal(model, config, gt_name, fx=1, fy=1, img_p=current_dir / "image.png", mask_p=current_dir / "mask.png"):
img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
2022-04-15 18:11:51 +02:00
res = model(img, mask, config)
2022-07-21 16:09:10 +02:00
cv2.imwrite(
2022-09-30 15:39:23 +02:00
str(save_dir / gt_name),
2022-07-21 16:09:10 +02:00
res,
[int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
)
2022-04-15 18:11:51 +02:00
"""
Note that JPEG is lossy compression, so even if it is the highest quality 100,
2022-07-14 10:49:03 +02:00
when the saved images is reloaded, a difference occurs with the original pixel value.
If you want to save the original images as it is, save it as PNG or BMP.
2022-04-15 18:11:51 +02:00
"""
2022-07-14 10:49:03 +02:00
# gt = cv2.imread(str(current_dir / gt_name), cv2.IMREAD_UNCHANGED)
# assert np.array_equal(res, gt)
2022-04-15 18:11:51 +02:00
2022-07-21 16:09:10 +02:00
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
2022-04-15 18:11:51 +02:00
def test_lama(strategy):
2022-09-05 07:07:25 +02:00
model = ModelManager(name="lama", device=device)
2022-07-21 16:09:10 +02:00
assert_equal(
model,
get_config(strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_result.png",
)
2022-04-15 18:11:51 +02:00
2022-07-21 16:09:10 +02:00
fx = 1.3
assert_equal(
model,
get_config(strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png",
fx=1.3,
)
2022-04-15 18:11:51 +02:00
2022-07-21 16:09:10 +02:00
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
@pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms])
2022-07-14 10:49:03 +02:00
def test_ldm(strategy, ldm_sampler):
2022-09-05 07:07:25 +02:00
model = ModelManager(name="ldm", device=device)
2022-07-14 10:49:03 +02:00
cfg = get_config(strategy, ldm_sampler=ldm_sampler)
2022-07-21 16:09:10 +02:00
assert_equal(
model, cfg, f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_result.png"
)
fx = 1.3
assert_equal(
model,
cfg,
f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_fx_{fx}_result.png",
fx=fx,
)
2022-07-14 10:49:03 +02:00
2022-07-21 16:09:10 +02:00
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
@pytest.mark.parametrize("zits_wireframe", [False, True])
2022-07-14 10:49:03 +02:00
def test_zits(strategy, zits_wireframe):
2022-09-05 07:07:25 +02:00
model = ModelManager(name="zits", device=device)
2022-07-14 10:49:03 +02:00
cfg = get_config(strategy, zits_wireframe=zits_wireframe)
# os.environ['ZITS_DEBUG_LINE_PATH'] = str(current_dir / 'zits_debug_line.jpg')
# os.environ['ZITS_DEBUG_EDGE_PATH'] = str(current_dir / 'zits_debug_edge.jpg')
2022-07-21 16:09:10 +02:00
assert_equal(
model,
cfg,
f"zits_{strategy[0].upper() + strategy[1:]}_wireframe_{zits_wireframe}_result.png",
)
fx = 1.3
assert_equal(
model,
cfg,
2022-08-22 17:24:02 +02:00
f"zits_{strategy.capitalize()}_wireframe_{zits_wireframe}_fx_{fx}_result.png",
2022-07-21 16:09:10 +02:00
fx=fx,
)
2022-08-22 17:24:02 +02:00
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL]
)
def test_mat(strategy):
2022-09-05 07:07:25 +02:00
model = ModelManager(name="mat", device=device)
2022-08-22 17:24:02 +02:00
cfg = get_config(strategy)
assert_equal(
model,
cfg,
f"mat_{strategy.capitalize()}_result.png",
)
2022-09-02 04:37:30 +02:00
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL]
)
def test_fcf(strategy):
2022-09-05 07:07:25 +02:00
model = ModelManager(name="fcf", device=device)
2022-09-02 04:37:30 +02:00
cfg = get_config(strategy)
assert_equal(
model,
cfg,
f"fcf_{strategy.capitalize()}_result.png",
fx=2,
fy=2
)
2022-09-04 10:00:42 +02:00
assert_equal(
model,
cfg,
f"fcf_{strategy.capitalize()}_result.png",
fx=3.8,
fy=2
)
2022-09-15 16:21:27 +02:00
2022-09-22 06:38:32 +02:00
2022-09-15 16:21:27 +02:00
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
2022-09-22 06:38:32 +02:00
@pytest.mark.parametrize("sampler", [SDSampler.ddim, SDSampler.pndm])
2022-09-29 15:56:33 +02:00
def test_sd(strategy, sampler):
2022-09-15 16:21:27 +02:00
def callback(step: int):
print(f"sd_step_{step}")
2022-09-22 06:38:32 +02:00
sd_steps = 50
2022-09-30 15:39:23 +02:00
model = ModelManager(name="sd1.4",
device=device,
hf_access_token=os.environ['HF_ACCESS_TOKEN'],
sd_run_local=False,
sd_disable_nsfw=False,
sd_cpu_textencoder=False,
2022-09-22 06:38:32 +02:00
callbacks=[callback])
2022-09-15 16:21:27 +02:00
cfg = get_config(strategy, prompt='a cat sitting on a bench', sd_steps=sd_steps)
2022-09-22 06:38:32 +02:00
cfg.sd_sampler = sampler
2022-09-15 16:21:27 +02:00
assert_equal(
model,
cfg,
2022-09-22 06:38:32 +02:00
f"sd_{strategy.capitalize()}_{sampler}_result.png",
2022-09-15 16:21:27 +02:00
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
2022-09-22 07:29:06 +02:00
assert_equal(
model,
cfg,
f"sd_{strategy.capitalize()}_{sampler}_blur_mask_result.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask_blur.png",
)
2022-09-29 15:56:33 +02:00
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
2022-09-30 15:39:23 +02:00
@pytest.mark.parametrize("disable_nsfw", [True, False])
def test_sd_run_local(strategy, sampler, disable_nsfw):
2022-09-29 15:56:33 +02:00
def callback(step: int):
print(f"sd_step_{step}")
sd_steps = 50
model = ModelManager(
name="sd1.4",
device=device,
2022-09-30 15:39:23 +02:00
# hf_access_token=os.environ.get('HF_ACCESS_TOKEN', None),
2022-09-29 15:56:33 +02:00
hf_access_token=None,
sd_run_local=True,
2022-09-30 15:39:23 +02:00
sd_disable_nsfw=disable_nsfw,
2022-09-29 15:56:33 +02:00
sd_cpu_textencoder=True,
)
cfg = get_config(strategy, prompt='a cat sitting on a bench', sd_steps=sd_steps)
cfg.sd_sampler = sampler
assert_equal(
model,
cfg,
f"sd_{strategy.capitalize()}_{sampler}_local_result.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
assert_equal(
model,
cfg,
f"sd_{strategy.capitalize()}_{sampler}_blur_mask_local_result.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask_blur.png",
)
2022-09-30 15:39:23 +02:00