add adjust mask feature

This commit is contained in:
Qing
2024-01-05 14:57:30 +08:00
parent 2996544e75
commit e889e527ab
18 changed files with 507 additions and 76 deletions

View File

@@ -0,0 +1,15 @@
import cv2
from lama_cleaner.helper import adjust_mask
from lama_cleaner.tests.utils import current_dir, save_dir
mask_p = current_dir / "overture-creations-5sI6fQgYIuo_mask.png"
def test_adjust_mask():
mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
res_mask = adjust_mask(mask, 0, "expand")
cv2.imwrite(str(save_dir / "adjust_mask_original.png"), res_mask)
res_mask = adjust_mask(mask, 40, "expand")
cv2.imwrite(str(save_dir / "adjust_mask_expand.png"), res_mask)
res_mask = adjust_mask(mask, 20, "shrink")
cv2.imwrite(str(save_dir / "adjust_mask_shrink.png"), res_mask)

View File

@@ -31,8 +31,6 @@ def assert_equal(
mask_p=current_dir / "mask.png",
):
img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
config.image = encode_pil_to_base64(Image.fromarray(img), 95, {})[0]
config.mask = encode_pil_to_base64(Image.fromarray(mask), 95, {})[0]
print(f"Input image shape: {img.shape}")
res = model(img, mask, config)
ok = cv2.imwrite(