IOPaint/iopaint/tests/test_plugins.py

122 lines
3.6 KiB
Python
Raw Normal View History

2023-04-06 15:55:20 +02:00
import hashlib
import os
import time
2024-01-02 04:07:35 +01:00
from PIL import Image
2023-04-06 15:55:20 +02:00
2024-01-05 08:19:23 +01:00
from iopaint.helper import encode_pil_to_base64, gen_frontend_mask
from iopaint.plugins.anime_seg import AnimeSeg
2024-02-20 02:03:11 +01:00
from iopaint.schema import RunPluginRequest, RemoveBGModel, InteractiveSegModel
2024-01-05 08:19:23 +01:00
from iopaint.tests.utils import check_device, current_dir, save_dir
2023-05-09 13:07:12 +02:00
2023-04-06 15:55:20 +02:00
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
2023-03-22 05:57:18 +01:00
import cv2
2023-03-26 14:42:31 +02:00
import pytest
2023-03-22 05:57:18 +01:00
2024-01-05 08:19:23 +01:00
from iopaint.plugins import (
2023-03-30 15:06:07 +02:00
RemoveBG,
RealESRGANUpscaler,
GFPGANPlugin,
RestoreFormerPlugin,
2023-04-06 15:55:20 +02:00
InteractiveSeg,
2023-03-30 15:06:07 +02:00
)
2023-03-22 05:57:18 +01:00
img_p = current_dir / "bunny.jpeg"
2023-04-06 15:55:20 +02:00
img_bytes = open(img_p, "rb").read()
2023-03-26 14:42:31 +02:00
bgr_img = cv2.imread(str(img_p))
rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
2024-01-02 04:07:35 +01:00
rgb_img_base64 = encode_pil_to_base64(Image.fromarray(rgb_img), 100, {})
bgr_img_base64 = encode_pil_to_base64(Image.fromarray(bgr_img), 100, {})
2023-03-26 14:42:31 +02:00
def _save(img, name):
cv2.imwrite(str(save_dir / name), img)
2023-03-22 05:57:18 +01:00
def test_remove_bg():
2024-02-08 09:49:54 +01:00
model = RemoveBG(RemoveBGModel.briaai_rmbg_1_4)
rgba_np_img = model.gen_image(
2024-01-02 04:07:35 +01:00
rgb_img, RunPluginRequest(name=RemoveBG.name, image=rgb_img_base64)
)
res = cv2.cvtColor(rgba_np_img, cv2.COLOR_RGBA2BGRA)
2023-03-26 14:42:31 +02:00
_save(res, "test_remove_bg.png")
bgr_np_img = model.gen_mask(
rgb_img, RunPluginRequest(name=RemoveBG.name, image=rgb_img_base64)
)
res_mask = gen_frontend_mask(bgr_np_img)
_save(res_mask, "test_remove_bg_frontend_mask.png")
assert len(bgr_np_img.shape) == 2
_save(bgr_np_img, "test_remove_bg_mask.jpeg")
2023-03-26 14:42:31 +02:00
2023-05-09 13:07:12 +02:00
def test_anime_seg():
model = AnimeSeg()
img = cv2.imread(str(current_dir / "anime_test.png"))
2024-01-02 04:07:35 +01:00
img_base64 = encode_pil_to_base64(Image.fromarray(img), 100, {})
res = model.gen_image(img, RunPluginRequest(name=AnimeSeg.name, image=img_base64))
2023-05-09 13:07:12 +02:00
assert len(res.shape) == 3
assert res.shape[-1] == 4
_save(res, "test_anime_seg.png")
res = model.gen_mask(img, RunPluginRequest(name=AnimeSeg.name, image=img_base64))
assert len(res.shape) == 2
_save(res, "test_anime_seg_mask.png")
2023-05-09 13:07:12 +02:00
2023-03-28 10:36:41 +02:00
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
2023-03-26 14:42:31 +02:00
def test_upscale(device):
2023-12-28 03:48:52 +01:00
check_device(device)
2023-03-26 14:42:31 +02:00
model = RealESRGANUpscaler("realesr-general-x4v3", device)
res = model.gen_image(
2024-01-02 04:07:35 +01:00
rgb_img,
RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=2),
)
2023-03-28 10:36:41 +02:00
_save(res, f"test_upscale_x2_{device}.png")
2023-03-22 05:57:18 +01:00
res = model.gen_image(
2024-01-02 04:07:35 +01:00
rgb_img,
RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=4),
)
2023-03-28 10:36:41 +02:00
_save(res, f"test_upscale_x4_{device}.png")
2023-03-22 05:57:18 +01:00
2023-03-28 10:36:41 +02:00
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
2023-03-26 14:42:31 +02:00
def test_gfpgan(device):
2023-12-28 03:48:52 +01:00
check_device(device)
2023-03-26 14:42:31 +02:00
model = GFPGANPlugin(device)
res = model.gen_image(
rgb_img, RunPluginRequest(name=GFPGANPlugin.name, image=rgb_img_base64)
)
2023-03-28 10:36:41 +02:00
_save(res, f"test_gfpgan_{device}.png")
2023-03-30 15:06:07 +02:00
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_restoreformer(device):
2023-12-28 03:48:52 +01:00
check_device(device)
2023-03-30 15:06:07 +02:00
model = RestoreFormerPlugin(device)
res = model.gen_image(
2024-01-02 04:07:35 +01:00
rgb_img, RunPluginRequest(name=RestoreFormerPlugin.name, image=rgb_img_base64)
)
2023-03-30 15:06:07 +02:00
_save(res, f"test_restoreformer_{device}.png")
2023-04-06 15:55:20 +02:00
2024-02-20 02:03:11 +01:00
@pytest.mark.parametrize("name", InteractiveSegModel.values())
2023-04-06 15:55:20 +02:00
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
2024-02-20 02:03:11 +01:00
def test_segment_anything(name, device):
2023-12-28 03:48:52 +01:00
check_device(device)
2024-02-20 02:03:11 +01:00
model = InteractiveSeg(name, device)
new_mask = model.gen_mask(
2024-01-02 04:07:35 +01:00
rgb_img,
RunPluginRequest(
name=InteractiveSeg.name,
image=rgb_img_base64,
clicks=([[448 // 2, 394 // 2, 1]]),
),
)
2023-04-06 15:55:20 +02:00
2024-02-20 02:03:11 +01:00
save_name = f"test_segment_anything_{name}_{device}.png"
2023-04-06 15:55:20 +02:00
_save(new_mask, save_name)