Add test codes

tests/test_config.py

@@ -0,0 +1,46 @@
+from PIL import Image, ImageChops
+import requests
+import os
+import torch
+import functools
+import numpy as np
+
+from modules.utils.paths import *
+
+
+TEST_IMAGE_URL = "https://github.com/microsoft/onnxjs-demo/raw/master/src/assets/EmotionSampleImages/sad_baby.jpg"
+TEST_IMAGE_PATH = os.path.join(PROJECT_ROOT_DIR, "tests", "test.png")
+TEST_EXPRESSION_OUTPUT_PATH = os.path.join(PROJECT_ROOT_DIR, "tests", "edited_expression.png")
+TEST_EXPRESSION_AAA = 100
+
+
+def download_image(url, path):
+    if os.path.exists(path):
+       return
+
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        with open(path, 'wb') as file:
+            for chunk in response.iter_content(1024):
+                file.write(chunk)
+        print(f"Image successfully downloaded to {path}")
+    else:
+        raise Exception(f"Failed to download image. Status code: {response.status_code}")
+
+
+def are_images_different(image1_path: str, image2_path: str):
+    image1 = Image.open(image1_path)
+    image2 = Image.open(image2_path)
+
+    diff = ImageChops.difference(image1, image2)
+
+    if diff.getbbox() is None:
+        return False
+    else:
+        return True
+
+
+@functools.lru_cache
+def is_cuda_available():
+    return torch.cuda.is_available()
+

tests/test_expression_editing.py

@@ -0,0 +1,34 @@
+import os
+import pytest
+
+from test_config import *
+from modules.live_portrait.live_portrait_inferencer import LivePortraitInferencer
+from modules.utils.image_helper import save_image
+
+
+@pytest.mark.parametrize(
+    "input_image,aaa",
+    [
+        (TEST_IMAGE_PATH, TEST_EXPRESSION_AAA),
+    ]
+)
+def test_expression_editing(
+    input_image: str,
+    aaa: int
+):
+    if not os.path.exists(TEST_IMAGE_PATH):
+        download_image(
+            TEST_IMAGE_URL,
+            TEST_IMAGE_PATH
+        )
+
+    inferencer = LivePortraitInferencer()
+
+    edited_expression = inferencer.edit_expression(
+        src_image=input_image,
+        aaa=aaa
+    )
+    save_image(numpy_array=edited_expression, output_path=TEST_EXPRESSION_OUTPUT_PATH)
+
+    assert os.path.exists(TEST_EXPRESSION_OUTPUT_PATH)
+    assert are_images_different(input_image, TEST_EXPRESSION_OUTPUT_PATH)