28 lines
717 B
Python
28 lines
717 B
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
from lama_cleaner.plugins import RemoveBG, RealESRGANUpscaler
|
||
|
|
|
||
|
|
current_dir = Path(__file__).parent.absolute().resolve()
|
||
|
|
save_dir = current_dir / "result"
|
||
|
|
save_dir.mkdir(exist_ok=True, parents=True)
|
||
|
|
img_p = current_dir / "bunny.jpeg"
|
||
|
|
|
||
|
|
|
||
|
|
def test_remove_bg():
|
||
|
|
model = RemoveBG()
|
||
|
|
img = cv2.imread(str(img_p))
|
||
|
|
res = model.forward(img)
|
||
|
|
cv2.imwrite(str(save_dir / "test_remove_bg.png"), res)
|
||
|
|
|
||
|
|
|
||
|
|
def test_upscale():
|
||
|
|
model = RealESRGANUpscaler("cpu")
|
||
|
|
img = cv2.imread(str(img_p))
|
||
|
|
res = model.forward(img, 2)
|
||
|
|
cv2.imwrite(str(save_dir / "test_upscale_x2.png"), res)
|
||
|
|
|
||
|
|
res = model.forward(img, 4)
|
||
|
|
cv2.imwrite(str(save_dir / "test_upscale_x4.png"), res)
|