fixed #627 env vars lost in Popen processes

tests/test_upscaler.py

@@ -0,0 +1,55 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+import ctypes
+import multiprocessing
+import os
+from pathlib import Path
+
+import utils
+from PIL import Image
+
+from video2x import Upscaler, Video2X
+
+
+def test_upscaling():
+    video2x = Video2X()
+    output_path = Path("data/test_video_output.mp4")
+    video2x.upscale(
+        Path("data/test_video.mp4"),
+        output_path,
+        None,
+        720,
+        3,
+        5,
+        0,
+        "waifu2x",
+    )
+    output_path.unlink()
+
+
+def test_upscale_image():
+
+    # initialize upscaler instance
+    processing_queue = multiprocessing.Queue(maxsize=30)
+    processed_frames = multiprocessing.Manager().list([None])
+    pause = multiprocessing.Value(ctypes.c_bool, False)
+    upscaler = Upscaler(processing_queue, processed_frames, pause)
+
+    image = Image.open("data/test_image.png")
+    upscaled_image = upscaler.upscale_image(image, 1680, 960, "waifu2x", 3)
+
+    reference_image = Image.open("data/test_image_ref.png")
+    assert utils.get_image_diff(upscaled_image, reference_image) < 0.5
+
+
+def test_get_scaling_tasks():
+    dimensions = [320, 240, 3840, 2160]
+
+    for algorithm, correct_answer in [
+        ("waifu2x", [2, 2, 2, 2]),
+        ["srmd", [3, 4]],
+        ("realsr", [4, 4]),
+        ("realcugan", [3, 4]),
+    ]:
+        assert Upscaler._get_scaling_tasks(*dimensions, algorithm) == correct_answer

tests/utils.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+from PIL import Image, ImageChops, ImageStat
+
+
+def get_image_diff(image0: Image.Image, image1: Image.Image) -> float:
+    """
+    calculate the percentage of differences between two images
+
+    :param image0 Image.Image: the first frame
+    :param image1 Image.Image: the second frame
+    :rtype float: the percent difference between the two images
+    """
+    difference = ImageChops.difference(image0, image1)
+    difference_stat = ImageStat.Stat(difference)
+    percent_diff = sum(difference_stat.mean) / (len(difference_stat.mean) * 255) * 100
+    return percent_diff

video2x/decoder.py

@@ -92,7 +92,7 @@ class VideoDecoder(threading.Thread):
                 ),
                 overwrite_output=True,
             ),
-            env={"AV_LOG_FORCE_COLOR": "TRUE"},
+            env=dict(AV_LOG_FORCE_COLOR="TRUE", **os.environ),
             stdin=subprocess.DEVNULL,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,

video2x/encoder.py

@@ -129,7 +129,7 @@ class VideoEncoder(threading.Thread):
                 ),
                 overwrite_output=True,
             ),
-            env={"AV_LOG_FORCE_COLOR": "TRUE"},
+            env=dict(AV_LOG_FORCE_COLOR="TRUE", **os.environ),
             stdin=subprocess.PIPE,
             stderr=subprocess.PIPE,
         )