File size: 1,398 Bytes
2fe55e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from pathlib import Path
import cv2
import numpy as np
from lama_cleaner.plugins import InteractiveSeg, Click
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
img_p = current_dir / "overture-creations-5sI6fQgYIuo.png"
def test_interactive_seg():
interactive_seg_model = InteractiveSeg()
img = cv2.imread(str(img_p))
pred = interactive_seg_model.forward(
img, clicks=[Click(coords=(256, 256), indx=0, is_positive=True)]
)
cv2.imwrite(str(save_dir / "test_interactive_seg.png"), pred)
def test_interactive_seg_with_negative_click():
interactive_seg_model = InteractiveSeg()
img = cv2.imread(str(img_p))
pred = interactive_seg_model.forward(
img,
clicks=[
Click(coords=(256, 256), indx=0, is_positive=True),
Click(coords=(384, 256), indx=1, is_positive=False),
],
)
cv2.imwrite(str(save_dir / "test_interactive_seg_negative.png"), pred)
def test_interactive_seg_with_prev_mask():
interactive_seg_model = InteractiveSeg()
img = cv2.imread(str(img_p))
mask = np.zeros_like(img)[:, :, 0]
pred = interactive_seg_model.forward(
img, clicks=[Click(coords=(256, 256), indx=0, is_positive=True)], prev_mask=mask
)
cv2.imwrite(str(save_dir / "test_interactive_seg_with_mask.png"), pred)
|