add plugins

lama_cleaner/tests/test_plugins.py

@@ -0,0 +1,27 @@
+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)