|
import hashlib |
|
import os |
|
import time |
|
from PIL import Image |
|
|
|
from iopaint.helper import encode_pil_to_base64, gen_frontend_mask |
|
from iopaint.plugins.anime_seg import AnimeSeg |
|
from iopaint.schema import RunPluginRequest, RemoveBGModel |
|
from iopaint.tests.utils import check_device, current_dir, save_dir |
|
|
|
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" |
|
|
|
import cv2 |
|
import pytest |
|
|
|
from iopaint.plugins import ( |
|
RemoveBG, |
|
RealESRGANUpscaler, |
|
GFPGANPlugin, |
|
RestoreFormerPlugin, |
|
InteractiveSeg, |
|
) |
|
|
|
img_p = current_dir / "bunny.jpeg" |
|
img_bytes = open(img_p, "rb").read() |
|
bgr_img = cv2.imread(str(img_p)) |
|
rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB) |
|
rgb_img_base64 = encode_pil_to_base64(Image.fromarray(rgb_img), 100, {}) |
|
bgr_img_base64 = encode_pil_to_base64(Image.fromarray(bgr_img), 100, {}) |
|
|
|
|
|
def _save(img, name): |
|
cv2.imwrite(str(save_dir / name), img) |
|
|
|
|
|
def test_remove_bg(): |
|
model = RemoveBG(RemoveBGModel.briaai_rmbg_1_4) |
|
rgba_np_img = model.gen_image( |
|
rgb_img, RunPluginRequest(name=RemoveBG.name, image=rgb_img_base64) |
|
) |
|
res = cv2.cvtColor(rgba_np_img, cv2.COLOR_RGBA2BGRA) |
|
_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") |
|
|
|
|
|
def test_anime_seg(): |
|
model = AnimeSeg() |
|
img = cv2.imread(str(current_dir / "anime_test.png")) |
|
img_base64 = encode_pil_to_base64(Image.fromarray(img), 100, {}) |
|
res = model.gen_image(img, RunPluginRequest(name=AnimeSeg.name, image=img_base64)) |
|
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") |
|
|
|
|
|
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) |
|
def test_upscale(device): |
|
check_device(device) |
|
model = RealESRGANUpscaler("realesr-general-x4v3", device) |
|
res = model.gen_image( |
|
rgb_img, |
|
RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=2), |
|
) |
|
_save(res, f"test_upscale_x2_{device}.png") |
|
|
|
res = model.gen_image( |
|
rgb_img, |
|
RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=4), |
|
) |
|
_save(res, f"test_upscale_x4_{device}.png") |
|
|
|
|
|
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) |
|
def test_gfpgan(device): |
|
check_device(device) |
|
model = GFPGANPlugin(device) |
|
res = model.gen_image( |
|
rgb_img, RunPluginRequest(name=GFPGANPlugin.name, image=rgb_img_base64) |
|
) |
|
_save(res, f"test_gfpgan_{device}.png") |
|
|
|
|
|
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) |
|
def test_restoreformer(device): |
|
check_device(device) |
|
model = RestoreFormerPlugin(device) |
|
res = model.gen_image( |
|
rgb_img, RunPluginRequest(name=RestoreFormerPlugin.name, image=rgb_img_base64) |
|
) |
|
_save(res, f"test_restoreformer_{device}.png") |
|
|
|
|
|
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) |
|
def test_segment_anything(device): |
|
check_device(device) |
|
model = InteractiveSeg("vit_l", device) |
|
new_mask = model.gen_mask( |
|
rgb_img, |
|
RunPluginRequest( |
|
name=InteractiveSeg.name, |
|
image=rgb_img_base64, |
|
clicks=([[448 // 2, 394 // 2, 1]]), |
|
), |
|
) |
|
|
|
save_name = f"test_segment_anything_{device}.png" |
|
_save(new_mask, save_name) |
|
|