Revert "added rich console patch for Google Colab"

This reverts commit db0b87597d.

Dockerfile

@@ -1,14 +1,16 @@
 # Name: Video2X Dockerfile
 # Creator: K4YT3X
 # Date Created: February 3, 2022
-# Last Modified: March 28, 2022
+# Last Modified: August 28, 2022
 
 # stage 1: build the python components into wheels
-FROM docker.io/nvidia/vulkan:1.2.133-450 AS builder
+FROM docker.io/nvidia/vulkan:1.3-470 AS builder
 ENV DEBIAN_FRONTEND=noninteractive
 
 COPY . /video2x
 WORKDIR /video2x
+RUN gpg --keyserver=keyserver.ubuntu.com --receive-keys A4B469963BF863CC \
+    && gpg --export A4B469963BF863CC > /etc/apt/trusted.gpg.d/cuda.gpg
 RUN apt-get update \
     && apt-get install -y --no-install-recommends \
         python3.8 python3-pip python3-opencv python3-pil \
@@ -17,7 +19,7 @@ RUN apt-get update \
     && pip wheel -w /wheels wheel pdm-pep517 .
 
 # stage 2: install wheels into the final image
-FROM docker.io/nvidia/vulkan:1.2.133-450
+FROM docker.io/nvidia/vulkan:1.3-470
 LABEL maintainer="K4YT3X <i@k4yt3x.com>" \
       org.opencontainers.image.source="https://github.com/k4yt3x/video2x" \
       org.opencontainers.image.description="A lossless video/GIF/image upscaler"

NOTICE

@@ -3,7 +3,7 @@ Copyright (c) 2018-2022 K4YT3X and contributors.
 
 This product depends on FFmpeg, which is available under
 the GNU Lesser General Public License 2.1. The source code can be found at
-https://www.gnu.org/licenses/agpl-3.0.txt.
+https://github.com/FFmpeg/FFmpeg.
 
 This product depends on waifu2x-ncnn-vulkan, which is available under
 the MIT License. The source code can be found at
@@ -41,3 +41,7 @@ https://github.com/python-pillow/Pillow.
 This product depends on Rich, which is available under
 the MIT License. The source code can be found at
 https://github.com/Textualize/rich.
+
+This product depends on pynput, which is available under
+the GNU Lesser General Public License 3.0. The source code can be found at
+https://github.com/moses-palmer/pynput.

README.md

@@ -82,18 +82,20 @@ Copyright (c) 2018-2022 K4YT3X and contributors.
 
 This project includes or depends on these following projects:
 
-| Project                                                             | License         |
+| Project                                                                 | License         |
-| ------------------------------------------------------------------- | --------------- |
+| ----------------------------------------------------------------------- | --------------- |
-| [FFmpeg](https://www.ffmpeg.org/)                                   | LGPLv2.1, GPLv2 |
+| [FFmpeg](https://www.ffmpeg.org/)                                       | LGPLv2.1, GPLv2 |
-| [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-ncnn-vulkan) | MIT License     |
+| [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-ncnn-vulkan)     | MIT License     |
-| [srmd-ncnn-vulkan](https://github.com/nihui/srmd-ncnn-vulkan)       | MIT License     |
+| [srmd-ncnn-vulkan](https://github.com/nihui/srmd-ncnn-vulkan)           | MIT License     |
-| [realsr-ncnn-vulkan](https://github.com/nihui/realsr-ncnn-vulkan)   | MIT License     |
+| [realsr-ncnn-vulkan](https://github.com/nihui/realsr-ncnn-vulkan)       | MIT License     |
-| [rife-ncnn-vulkan](https://github.com/nihui/rife-ncnn-vulkan)       | MIT License     |
+| [rife-ncnn-vulkan](https://github.com/nihui/rife-ncnn-vulkan)           | MIT License     |
-| [ffmpeg-python](https://github.com/kkroening/ffmpeg-python)         | Apache-2.0      |
+| [realcugan-ncnn-vulkan](https://github.com/nihui/realcugan-ncnn-vulkan) | MIT License     |
-| [Loguru](https://github.com/Delgan/loguru)                          | MIT License     |
+| [ffmpeg-python](https://github.com/kkroening/ffmpeg-python)             | Apache-2.0      |
-| [opencv-python](https://github.com/opencv/opencv-python)            | MIT License     |
+| [Loguru](https://github.com/Delgan/loguru)                              | MIT License     |
-| [Pillow](https://github.com/python-pillow/Pillow)                   | HPND License    |
+| [opencv-python](https://github.com/opencv/opencv-python)                | MIT License     |
-| [Rich](https://github.com/Textualize/rich)                          | MIT License     |
+| [Pillow](https://github.com/python-pillow/Pillow)                       | HPND License    |
+| [Rich](https://github.com/Textualize/rich)                              | MIT License     |
+| [pynput](https://github.com/moses-palmer/pynput)                        | LGPLv3.0        |
 
 Legacy versions of this project includes or depends on these following projects:
 
@@ -117,6 +119,7 @@ Appreciations given to the following personnel who have contributed significantl
 - [@ddouglas87](https://github.com/ddouglas87)
 - [@lhanjian](https://github.com/lhanjian)
 - [@ArchieMeng](https://github.com/archiemeng)
+- [@nihui](https://github.com/nihui)
 
 ## Similar Projects
 

licenses/LICENSE-realcugan-ncnn-vulkan

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2019 nihui
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

licenses/lgpl-3.0.txt

@@ -0,0 +1,165 @@
+                   GNU LESSER GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+
+  This version of the GNU Lesser General Public License incorporates
+the terms and conditions of version 3 of the GNU General Public
+License, supplemented by the additional permissions listed below.
+
+  0. Additional Definitions.
+
+  As used herein, "this License" refers to version 3 of the GNU Lesser
+General Public License, and the "GNU GPL" refers to version 3 of the GNU
+General Public License.
+
+  "The Library" refers to a covered work governed by this License,
+other than an Application or a Combined Work as defined below.
+
+  An "Application" is any work that makes use of an interface provided
+by the Library, but which is not otherwise based on the Library.
+Defining a subclass of a class defined by the Library is deemed a mode
+of using an interface provided by the Library.
+
+  A "Combined Work" is a work produced by combining or linking an
+Application with the Library.  The particular version of the Library
+with which the Combined Work was made is also called the "Linked
+Version".
+
+  The "Minimal Corresponding Source" for a Combined Work means the
+Corresponding Source for the Combined Work, excluding any source code
+for portions of the Combined Work that, considered in isolation, are
+based on the Application, and not on the Linked Version.
+
+  The "Corresponding Application Code" for a Combined Work means the
+object code and/or source code for the Application, including any data
+and utility programs needed for reproducing the Combined Work from the
+Application, but excluding the System Libraries of the Combined Work.
+
+  1. Exception to Section 3 of the GNU GPL.
+
+  You may convey a covered work under sections 3 and 4 of this License
+without being bound by section 3 of the GNU GPL.
+
+  2. Conveying Modified Versions.
+
+  If you modify a copy of the Library, and, in your modifications, a
+facility refers to a function or data to be supplied by an Application
+that uses the facility (other than as an argument passed when the
+facility is invoked), then you may convey a copy of the modified
+version:
+
+   a) under this License, provided that you make a good faith effort to
+   ensure that, in the event an Application does not supply the
+   function or data, the facility still operates, and performs
+   whatever part of its purpose remains meaningful, or
+
+   b) under the GNU GPL, with none of the additional permissions of
+   this License applicable to that copy.
+
+  3. Object Code Incorporating Material from Library Header Files.
+
+  The object code form of an Application may incorporate material from
+a header file that is part of the Library.  You may convey such object
+code under terms of your choice, provided that, if the incorporated
+material is not limited to numerical parameters, data structure
+layouts and accessors, or small macros, inline functions and templates
+(ten or fewer lines in length), you do both of the following:
+
+   a) Give prominent notice with each copy of the object code that the
+   Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the object code with a copy of the GNU GPL and this license
+   document.
+
+  4. Combined Works.
+
+  You may convey a Combined Work under terms of your choice that,
+taken together, effectively do not restrict modification of the
+portions of the Library contained in the Combined Work and reverse
+engineering for debugging such modifications, if you also do each of
+the following:
+
+   a) Give prominent notice with each copy of the Combined Work that
+   the Library is used in it and that the Library and its use are
+   covered by this License.
+
+   b) Accompany the Combined Work with a copy of the GNU GPL and this license
+   document.
+
+   c) For a Combined Work that displays copyright notices during
+   execution, include the copyright notice for the Library among
+   these notices, as well as a reference directing the user to the
+   copies of the GNU GPL and this license document.
+
+   d) Do one of the following:
+
+       0) Convey the Minimal Corresponding Source under the terms of this
+       License, and the Corresponding Application Code in a form
+       suitable for, and under terms that permit, the user to
+       recombine or relink the Application with a modified version of
+       the Linked Version to produce a modified Combined Work, in the
+       manner specified by section 6 of the GNU GPL for conveying
+       Corresponding Source.
+
+       1) Use a suitable shared library mechanism for linking with the
+       Library.  A suitable mechanism is one that (a) uses at run time
+       a copy of the Library already present on the user's computer
+       system, and (b) will operate properly with a modified version
+       of the Library that is interface-compatible with the Linked
+       Version.
+
+   e) Provide Installation Information, but only if you would otherwise
+   be required to provide such information under section 6 of the
+   GNU GPL, and only to the extent that such information is
+   necessary to install and execute a modified version of the
+   Combined Work produced by recombining or relinking the
+   Application with a modified version of the Linked Version. (If
+   you use option 4d0, the Installation Information must accompany
+   the Minimal Corresponding Source and Corresponding Application
+   Code. If you use option 4d1, you must provide the Installation
+   Information in the manner specified by section 6 of the GNU GPL
+   for conveying Corresponding Source.)
+
+  5. Combined Libraries.
+
+  You may place library facilities that are a work based on the
+Library side by side in a single library together with other library
+facilities that are not Applications and are not covered by this
+License, and convey such a combined library under terms of your
+choice, if you do both of the following:
+
+   a) Accompany the combined library with a copy of the same work based
+   on the Library, uncombined with any other library facilities,
+   conveyed under the terms of this License.
+
+   b) Give prominent notice with the combined library that part of it
+   is a work based on the Library, and explaining where to find the
+   accompanying uncombined form of the same work.
+
+  6. Revised Versions of the GNU Lesser General Public License.
+
+  The Free Software Foundation may publish revised and/or new versions
+of the GNU Lesser General Public License from time to time. Such new
+versions will be similar in spirit to the present version, but may
+differ in detail to address new problems or concerns.
+
+  Each version is given a distinguishing version number. If the
+Library as you received it specifies that a certain numbered version
+of the GNU Lesser General Public License "or any later version"
+applies to it, you have the option of following the terms and
+conditions either of that published version or of any later version
+published by the Free Software Foundation. If the Library as you
+received it does not specify a version number of the GNU Lesser
+General Public License, you may choose any version of the GNU Lesser
+General Public License ever published by the Free Software Foundation.
+
+  If the Library as you received it specifies that a proxy can decide
+whether future versions of the GNU Lesser General Public License shall
+apply, that proxy's public statement of acceptance of any version is
+permanent authorization for you to choose that version for the
+Library.

pyproject.toml

@@ -3,7 +3,7 @@ name = "video2x"
 description = "A video/image upscaling and frame interpolation framework"
 readme = "README.md"
 requires-python = ">=3.7"
-license = { text = "AGPL-3.0-or-later" }
+license-expression = "AGPL-3.0-or-later"
 keywords = [
   "super-resolution",
   "upscaling",
@@ -27,15 +27,15 @@ classifiers = [
 dependencies = [
   "ffmpeg-python>=0.2.0",
   "loguru>=0.6.0",
-  "opencv-python>=4.5.5.64",
+  "opencv-python==4.5.5.64",
-  "pillow>=9.0.1",
+  "pillow>=9.1.0",
+  "pynput>=1.7.6",
   "rich>=12.0.0",
   "waifu2x-ncnn-vulkan-python>=1.0.2rc3",
   "srmd-ncnn-vulkan-python>=1.0.2",
   "realsr-ncnn-vulkan-python>=1.0.4",
   "rife-ncnn-vulkan-python>=1.1.2.post3",
-  "realcugan-ncnn-vulkan-python>=1.0.0",
+  "realcugan-ncnn-vulkan-python>=1.0.2",
-  "pynput>=1.7.6",
 ]
 dynamic = ["version"]
 
@@ -48,9 +48,12 @@ changelog = "https://github.com/k4yt3x/video2x/releases"
 [project.scripts]
 video2x = "video2x:main"
 
+[tool.isort]
+profile = "black"
+
 [tool.pdm]
 version = { from = "video2x/__init__.py" }
 
 [build-system]
-requires = ["pdm-pep517"]
+requires = ["pdm-pep517>=0.12.0"]
 build-backend = "pdm.pep517.api"

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/__init__.py

@@ -19,13 +19,16 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
 Name: Package Init
 Author: K4YT3X
 Date Created: July 3, 2021
-Last Modified: February 16, 2022
+Last Modified: August 28, 2022
 """
 
 # version assignment has to precede imports to
 # prevent setup.cfg from producing import errors
-__version__ = "5.0.0-beta5"
+__version__ = "5.0.0-beta6"
 
+# flake8: noqa
+# let flake8 ignore this file to avoid F401 warnings
+# generated by the following lines
 from .interpolator import Interpolator
 from .upscaler import Upscaler
 from .video2x import Video2X

video2x/__main__.py

@@ -22,9 +22,218 @@ Date Created: July 3, 2021
 Last Modified: February 26, 2022
 """
 
+import argparse
+import os
+import pathlib
 import sys
 
-from .video2x import main
+from loguru import logger
+from rich import print as rich_print
+
+from . import __version__
+from .video2x import LOGURU_FORMAT, Video2X
+
+LEGAL_INFO = f"""Video2X\t\t{__version__}
+Author:\t\tK4YT3X
+License:\tGNU AGPL v3
+Github Page:\thttps://github.com/k4yt3x/video2x
+Contact:\ti@k4yt3x.com"""
+
+# algorithms available for upscaling tasks
+UPSCALING_ALGORITHMS = [
+    "waifu2x",
+    "srmd",
+    "realsr",
+    "realcugan",
+]
+
+# algorithms available for frame interpolation tasks
+INTERPOLATION_ALGORITHMS = ["rife"]
+
+
+def parse_arguments() -> argparse.Namespace:
+    """
+    parse command line arguments
+
+    :rtype argparse.Namespace: command parsing results
+    """
+    parser = argparse.ArgumentParser(
+        prog="video2x",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--version", help="show version information and exit", action="store_true"
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=pathlib.Path,
+        help="input file/directory path",
+        required=True,
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=pathlib.Path,
+        help="output file/directory path",
+        required=True,
+    )
+    parser.add_argument(
+        "-p", "--processes", type=int, help="number of processes to launch", default=1
+    )
+    parser.add_argument(
+        "-l",
+        "--loglevel",
+        choices=["trace", "debug", "info", "success", "warning", "error", "critical"],
+        default="info",
+    )
+
+    # upscaler arguments
+    action = parser.add_subparsers(
+        help="action to perform", dest="action", required=True
+    )
+
+    upscale = action.add_parser(
+        "upscale",
+        help="upscale a file",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        add_help=False,
+    )
+    upscale.add_argument(
+        "--help", action="help", help="show this help message and exit"
+    )
+    upscale.add_argument("-w", "--width", type=int, help="output width")
+    upscale.add_argument("-h", "--height", type=int, help="output height")
+    upscale.add_argument("-n", "--noise", type=int, help="denoise level", default=3)
+    upscale.add_argument(
+        "-a",
+        "--algorithm",
+        choices=UPSCALING_ALGORITHMS,
+        help="algorithm to use for upscaling",
+        default=UPSCALING_ALGORITHMS[0],
+    )
+    upscale.add_argument(
+        "-t",
+        "--threshold",
+        type=float,
+        help=(
+            "skip if the percent difference between two adjacent frames is below this"
+            " value; set to 0 to process all frames"
+        ),
+        default=0,
+    )
+
+    # interpolator arguments
+    interpolate = action.add_parser(
+        "interpolate",
+        help="interpolate frames for file",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        add_help=False,
+    )
+    interpolate.add_argument(
+        "--help", action="help", help="show this help message and exit"
+    )
+    interpolate.add_argument(
+        "-a",
+        "--algorithm",
+        choices=UPSCALING_ALGORITHMS,
+        help="algorithm to use for upscaling",
+        default=INTERPOLATION_ALGORITHMS[0],
+    )
+    interpolate.add_argument(
+        "-t",
+        "--threshold",
+        type=float,
+        help=(
+            "skip if the percent difference between two adjacent frames exceeds this"
+            " value; set to 100 to interpolate all frames"
+        ),
+        default=10,
+    )
+
+    return parser.parse_args()
+
+
+def main() -> int:
+    """
+    command line entrypoint for direct CLI invocation
+
+    :rtype int: 0 if completed successfully, else other int
+    """
+
+    try:
+        # display version and lawful informaition
+        if "--version" in sys.argv:
+            rich_print(LEGAL_INFO)
+            return 0
+
+        # parse command line arguments
+        args = parse_arguments()
+
+        # check input/output file paths
+        if not args.input.exists():
+            logger.critical(f"Cannot find input file: {args.input}")
+            return 1
+        if not args.input.is_file():
+            logger.critical("Input path is not a file")
+            return 1
+        if not args.output.parent.exists():
+            logger.critical(f"Output directory does not exist: {args.output.parent}")
+            return 1
+
+        # set logger level
+        if os.environ.get("LOGURU_LEVEL") is None:
+            os.environ["LOGURU_LEVEL"] = args.loglevel.upper()
+
+        # remove default handler
+        logger.remove()
+
+        # add new sink with custom handler
+        logger.add(sys.stderr, colorize=True, format=LOGURU_FORMAT)
+
+        # print package version and copyright notice
+        logger.opt(colors=True).info(f"<magenta>Video2X {__version__}</magenta>")
+        logger.opt(colors=True).info(
+            "<magenta>Copyright (C) 2018-2022 K4YT3X and contributors.</magenta>"
+        )
+
+        # initialize video2x object
+        video2x = Video2X()
+
+        if args.action == "upscale":
+            video2x.upscale(
+                args.input,
+                args.output,
+                args.width,
+                args.height,
+                args.noise,
+                args.processes,
+                args.threshold,
+                args.algorithm,
+            )
+
+        elif args.action == "interpolate":
+            video2x.interpolate(
+                args.input,
+                args.output,
+                args.processes,
+                args.threshold,
+                args.algorithm,
+            )
+
+    # don't print the traceback for manual terminations
+    except KeyboardInterrupt:
+        return 2
+
+    except Exception as error:
+        logger.exception(error)
+        return 1
+
+    # if no exceptions were produced
+    else:
+        logger.success("Processing completed successfully")
+        return 0
+
 
 if __name__ == "__main__":
     sys.exit(main())

video2x/decoder.py

@@ -19,22 +19,19 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
 Name: Video Decoder
 Author: K4YT3X
 Date Created: June 17, 2021
-Last Modified: March 21, 2022
+Last Modified: April 9, 2022
 """
 
 import contextlib
-import multiprocessing
 import os
 import pathlib
-import queue
 import signal
 import subprocess
-import threading
+from multiprocessing import Queue
-import time
+from queue import Full
-from multiprocessing.sharedctypes import Synchronized
+from threading import Thread
 
 import ffmpeg
-from loguru import logger
 from PIL import Image
 
 from .pipe_printer import PipePrinter
@@ -51,36 +48,38 @@ LOGURU_FFMPEG_LOGLEVELS = {
 }
 
 
-class VideoDecoder(threading.Thread):
+class VideoDecoder:
+    """
+    A video decoder that generates frames read from FFmpeg.
+
+    :param input_path pathlib.Path: the input file's path
+    :param input_width int: the input file's width
+    :param input_height int: the input file's height
+    :param frame_rate float: the input file's frame rate
+    :param pil_ignore_max_image_pixels bool: setting this to True
+        disables PIL's "possible DDoS" warning
+    """
+
     def __init__(
         self,
         input_path: pathlib.Path,
         input_width: int,
         input_height: int,
         frame_rate: float,
-        processing_queue: multiprocessing.Queue,
+        pil_ignore_max_image_pixels: bool = True,
-        processing_settings: tuple,
-        pause: Synchronized,
-        ignore_max_image_pixels=True,
     ) -> None:
-        threading.Thread.__init__(self)
-        self.running = False
         self.input_path = input_path
         self.input_width = input_width
         self.input_height = input_height
-        self.processing_queue = processing_queue
-        self.processing_settings = processing_settings
-        self.pause = pause
 
         # this disables the "possible DDoS" warning
-        if ignore_max_image_pixels:
+        if pil_ignore_max_image_pixels is True:
             Image.MAX_IMAGE_PIXELS = None
 
-        self.exception = None
         self.decoder = subprocess.Popen(
             ffmpeg.compile(
                 ffmpeg.input(input_path, r=frame_rate)["v"]
-                .output("pipe:1", format="rawvideo", pix_fmt="rgb24", vsync="cfr")
+                .output("pipe:1", format="rawvideo", pix_fmt="rgb24")
                 .global_args("-hide_banner")
                 .global_args("-nostats")
                 .global_args("-nostdin")
@@ -92,7 +91,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,
@@ -102,83 +101,33 @@ class VideoDecoder(threading.Thread):
         self.pipe_printer = PipePrinter(self.decoder.stderr)
         self.pipe_printer.start()
 
-    def run(self) -> None:
+    def __iter__(self):
-        self.running = True
 
-        # the index of the frame
+        # continue yielding while FFmpeg continues to produce output
-        frame_index = 0
+        # it is possible to use := for this block to be more concise
+        # but it is purposefully avoided to remain compatible with Python 3.7
+        buffer = self.decoder.stdout.read(3 * self.input_width * self.input_height)
 
-        # create placeholder for previous frame
+        while len(buffer) > 0:
-        # used in interpolate mode
-        previous_image = None
 
-        # continue running until an exception occurs
+            # convert raw bytes into image object
-        # or all frames have been decoded
+            frame = Image.frombytes(
-        while self.running:
+                "RGB", (self.input_width, self.input_height), buffer
+            )
 
-            # pause if pause flag is set
+            # return this frame
-            if self.pause.value is True:
+            yield frame
-                time.sleep(0.1)
-                continue
 
-            try:
+            # read the next frame
-                buffer = self.decoder.stdout.read(
+            buffer = self.decoder.stdout.read(3 * self.input_width * self.input_height)
-                    3 * self.input_width * self.input_height
-                )
 
-                # source depleted (decoding finished)
+        # automatically self-join and clean up after iterations are done
-                # after the last frame has been decoded
+        self.join()
-                # read will return nothing
-                if len(buffer) == 0:
-                    self.stop()
-                    continue
 
-                # convert raw bytes into image object
+    def kill(self):
-                image = Image.frombytes(
+        self.decoder.send_signal(signal.SIGKILL)
-                    "RGB", (self.input_width, self.input_height), buffer
-                )
 
-                # keep checking if the running flag is set to False
+    def join(self):
-                # while waiting to put the next image into the queue
-                while self.running:
-                    with contextlib.suppress(queue.Full):
-                        self.processing_queue.put(
-                            (
-                                frame_index,
-                                (previous_image, image),
-                                self.processing_settings,
-                            ),
-                            timeout=0.1,
-                        )
-                        break
-
-                previous_image = image
-                frame_index += 1
-
-            # most likely "not enough image data"
-            except ValueError as e:
-                self.exception = e
-
-                # ignore queue closed
-                if "is closed" not in str(e):
-                    logger.exception(e)
-                break
-
-            # send exceptions into the client connection pipe
-            except Exception as e:
-                self.exception = e
-                logger.exception(e)
-                break
-        else:
-            logger.debug("Decoding queue depleted")
-
-        # flush the remaining data in STDOUT and STDERR
-        self.decoder.stdout.flush()
-        self.decoder.stderr.flush()
-
-        # send SIGINT (2) to FFmpeg
-        # this instructs it to finalize and exit
-        self.decoder.send_signal(signal.SIGINT)
 
         # close PIPEs to prevent process from getting stuck
         self.decoder.stdout.close()
@@ -191,8 +140,38 @@ class VideoDecoder(threading.Thread):
         self.pipe_printer.stop()
         self.pipe_printer.join()
 
-        logger.info("Decoder thread exiting")
-        return super().run()
 
-    def stop(self) -> None:
+class VideoDecoderThread(Thread):
+    def __init__(
+        self, tasks_queue: Queue, decoder: VideoDecoder, processing_settings: tuple
+    ):
+        super().__init__()
+
+        self.tasks_queue = tasks_queue
+        self.decoder = decoder
+        self.processing_settings = processing_settings
+        self.running = False
+
+    def run(self):
+        self.running = True
+        previous_frame = None
+        for frame_index, frame in enumerate(self.decoder):
+
+            while True:
+
+                # check for the stop signal
+                if self.running is False:
+                    self.decoder.join()
+                    return
+
+                with contextlib.suppress(Full):
+                    self.tasks_queue.put(
+                        (frame_index, previous_frame, frame, self.processing_settings),
+                        timeout=0.1,
+                    )
+                    break
+
+            previous_frame = frame
+
+    def stop(self):
         self.running = False

video2x/encoder.py

@@ -19,20 +19,16 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
 Name: Video Encoder
 Author: K4YT3X
 Date Created: June 17, 2021
-Last Modified: March 20, 2022
+Last Modified: August 28, 2022
 """
 
 import os
 import pathlib
 import signal
 import subprocess
-import threading
-import time
-from multiprocessing.managers import ListProxy
-from multiprocessing.sharedctypes import Synchronized
 
 import ffmpeg
-from loguru import logger
+from PIL import Image
 
 from .pipe_printer import PipePrinter
 
@@ -48,7 +44,7 @@ LOGURU_FFMPEG_LOGLEVELS = {
 }
 
 
-class VideoEncoder(threading.Thread):
+class VideoEncoder:
     def __init__(
         self,
         input_path: pathlib.Path,
@@ -56,36 +52,20 @@ class VideoEncoder(threading.Thread):
         output_path: pathlib.Path,
         output_width: int,
         output_height: int,
-        total_frames: int,
-        processed_frames: ListProxy,
-        processed: Synchronized,
-        pause: Synchronized,
         copy_audio: bool = True,
         copy_subtitle: bool = True,
         copy_data: bool = False,
         copy_attachments: bool = False,
     ) -> None:
-        threading.Thread.__init__(self)
-        self.running = False
-        self.input_path = input_path
-        self.output_path = output_path
-        self.total_frames = total_frames
-        self.processed_frames = processed_frames
-        self.processed = processed
-        self.pause = pause
-
-        # stores exceptions if the thread exits with errors
-        self.exception = None
 
         # create FFmpeg input for the original input video
-        self.original = ffmpeg.input(input_path)
+        original = ffmpeg.input(input_path)
 
         # define frames as input
         frames = ffmpeg.input(
             "pipe:0",
             format="rawvideo",
             pix_fmt="rgb24",
-            vsync="cfr",
             s=f"{output_width}x{output_height}",
             r=frame_rate,
         )
@@ -93,11 +73,11 @@ class VideoEncoder(threading.Thread):
         # copy additional streams from original file
         # https://ffmpeg.org/ffmpeg.html#Stream-specifiers-1
         additional_streams = [
-            # self.original["1:v?"],
+            # original["1:v?"],
-            self.original["a?"] if copy_audio is True else None,
+            original["a?"] if copy_audio is True else None,
-            self.original["s?"] if copy_subtitle is True else None,
+            original["s?"] if copy_subtitle is True else None,
-            self.original["d?"] if copy_data is True else None,
+            original["d?"] if copy_data is True else None,
-            self.original["t?"] if copy_attachments is True else None,
+            original["t?"] if copy_attachments is True else None,
         ]
 
         # run FFmpeg and produce final output
@@ -106,9 +86,9 @@ class VideoEncoder(threading.Thread):
                 ffmpeg.output(
                     frames,
                     *[s for s in additional_streams if s is not None],
-                    str(self.output_path),
+                    str(output_path),
                     vcodec="libx264",
-                    vsync="cfr",
+                    scodec="copy",
                     pix_fmt="yuv420p",
                     crf=17,
                     preset="veryslow",
@@ -128,7 +108,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,
         )
@@ -137,49 +117,26 @@ class VideoEncoder(threading.Thread):
         self.pipe_printer = PipePrinter(self.encoder.stderr)
         self.pipe_printer.start()
 
-    def run(self) -> None:
+    def kill(self):
-        self.running = True
+        self.encoder.send_signal(signal.SIGKILL)
-        frame_index = 0
-        while self.running and frame_index < self.total_frames:
 
-            # pause if pause flag is set
+    def write(self, frame: Image.Image) -> None:
-            if self.pause.value is True:
+        """
-                time.sleep(0.1)
+        write a frame into FFmpeg encoder's STDIN
-                continue
 
-            try:
+        :param frame Image.Image: the Image object to use for writing
-                image = self.processed_frames[frame_index]
+        """
-                if image is None:
+        self.encoder.stdin.write(frame.tobytes())
-                    time.sleep(0.1)
-                    continue
-
-                # send the image to FFmpeg for encoding
-                self.encoder.stdin.write(image.tobytes())
-
-                # remove the image from memory
-                self.processed_frames[frame_index] = None
-
-                with self.processed.get_lock():
-                    self.processed.value += 1
-
-                frame_index += 1
-
-            # send exceptions into the client connection pipe
-            except Exception as e:
-                self.exception = e
-                logger.exception(e)
-                break
-        else:
-            logger.debug("Encoding queue depleted")
 
+    def join(self) -> None:
+        """
+        signal the encoder that all frames have been sent and the FFmpeg
+        should be instructed to wrap-up the processing
+        """
         # flush the remaining data in STDIN and STDERR
         self.encoder.stdin.flush()
         self.encoder.stderr.flush()
 
-        # send SIGINT (2) to FFmpeg
-        # this instructs it to finalize and exit
-        self.encoder.send_signal(signal.SIGINT)
-
         # close PIPEs to prevent process from getting stuck
         self.encoder.stdin.close()
         self.encoder.stderr.close()
@@ -190,9 +147,3 @@ class VideoEncoder(threading.Thread):
         # wait for PIPE printer to exit
         self.pipe_printer.stop()
         self.pipe_printer.join()
-
-        logger.info("Encoder thread exiting")
-        return super().run()
-
-    def stop(self) -> None:
-        self.running = False

video2x/interpolator.py

@@ -44,11 +44,13 @@ class Interpolator(multiprocessing.Process):
         pause: Synchronized,
     ) -> None:
         multiprocessing.Process.__init__(self)
-        self.running = False
         self.processing_queue = processing_queue
         self.processed_frames = processed_frames
         self.pause = pause
 
+        self.running = False
+        self.processor_objects = {}
+
         signal.signal(signal.SIGTERM, self._stop)
 
     def run(self) -> None:
@@ -56,8 +58,7 @@ class Interpolator(multiprocessing.Process):
         logger.opt(colors=True).info(
             f"Interpolator process <blue>{self.name}</blue> initiating"
         )
-        processor_objects = {}
+        while self.running is True:
-        while self.running:
             try:
                 # pause if pause flag is set
                 if self.pause.value is True:
@@ -80,6 +81,7 @@ class Interpolator(multiprocessing.Process):
                 if image0 is None:
                     continue
 
+                # calculate the %diff between the current frame and the previous frame
                 difference = ImageChops.difference(image0, image1)
                 difference_stat = ImageStat.Stat(difference)
                 difference_ratio = (
@@ -92,10 +94,10 @@ class Interpolator(multiprocessing.Process):
 
                     # select a processor object with the required settings
                     # create a new object if none are available
-                    processor_object = processor_objects.get(algorithm)
+                    processor_object = self.processor_objects.get(algorithm)
                     if processor_object is None:
                         processor_object = ALGORITHM_CLASSES[algorithm](0)
-                        processor_objects[algorithm] = processor_object
+                        self.processor_objects[algorithm] = processor_object
                     interpolated_image = processor_object.process(image0, image1)
 
                 # if the difference is greater than threshold
@@ -112,8 +114,8 @@ class Interpolator(multiprocessing.Process):
             except (SystemExit, KeyboardInterrupt):
                 break
 
-            except Exception as e:
+            except Exception as error:
-                logger.exception(e)
+                logger.exception(error)
                 break
 
         logger.opt(colors=True).info(

video2x/pipe_printer.py

@@ -47,7 +47,7 @@ class PipePrinter(threading.Thread):
         self.running = True
 
         # keep printing contents in the PIPE
-        while self.running:
+        while self.running is True:
             time.sleep(0.5)
 
             try:

video2x/processor.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Copyright (C) 2018-2022 K4YT3X and contributors.
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+Name: Processor Abstract Class
+Author: K4YT3X
+Date Created: April 9, 2022
+Last Modified: April 9, 2022
+"""
+
+from abc import ABC, abstractmethod
+from multiprocessing import Queue
+from multiprocessing.managers import DictProxy
+from multiprocessing.sharedctypes import Synchronized
+
+from PIL import Image, ImageChops, ImageStat
+
+
+class Processor(ABC):
+    def __init__(
+        self, tasks_queue: Queue, processed_frames: DictProxy, pause_flag: Synchronized
+    ) -> None:
+        self.tasks_queue = tasks_queue
+        self.processed_frames = processed_frames
+        self.pause_flag = pause_flag
+
+    @abstractmethod
+    def process(self):
+        raise NotImplementedError
+
+    @staticmethod
+    def get_image_diff(image0: Image.Image, image1: Image.Image) -> float:
+        """
+        get the percentage difference between two images
+
+        :param image0 Image.Image: the image to compare
+        :param image1 Image.Image: the image to compare against
+        :rtype float: precentage difference between two frames
+        """
+        difference_stat = ImageStat.Stat(ImageChops.difference(image0, image1))
+        return sum(difference_stat.mean) / (len(difference_stat.mean) * 255) * 100
+
+    """
+    def run(
+        self,
+    ) -> None:
+        self.running = True
+        while self.running is True:
+            self.process()
+        self.running = False
+        return super().run()
+
+    def stop(self, _signal_number, _frame) -> None:
+        self.running = False
+    """

video2x/upscaler.py

@@ -19,189 +19,184 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
 Name: Upscaler
 Author: K4YT3X
 Date Created: May 27, 2021
-Last Modified: March 20, 2022
+Last Modified: April 10, 2022
 """
 
 import math
-import multiprocessing
-import queue
-import signal
 import time
-from multiprocessing.managers import ListProxy
-from multiprocessing.sharedctypes import Synchronized
 
-from loguru import logger
+from PIL import Image
-from PIL import Image, ImageChops, ImageStat
 from realcugan_ncnn_vulkan_python import Realcugan
 from realsr_ncnn_vulkan_python import Realsr
 from srmd_ncnn_vulkan_python import Srmd
 from waifu2x_ncnn_vulkan_python import Waifu2x
 
-# fixed scaling ratios supported by the algorithms
+from .processor import Processor
-# that only support certain fixed scale ratios
-ALGORITHM_FIXED_SCALING_RATIOS = {
-    "waifu2x": [1, 2],
-    "srmd": [2, 3, 4],
-    "realsr": [4],
-    "realcugan": [1, 2, 3, 4],
-}
-
-ALGORITHM_CLASSES = {
-    "waifu2x": Waifu2x,
-    "srmd": Srmd,
-    "realsr": Realsr,
-    "realcugan": Realcugan,
-}
 
 
-class Upscaler(multiprocessing.Process):
+class Upscaler:
-    def __init__(
+    # fixed scaling ratios supported by the algorithms
-        self,
+    # that only support certain fixed scale ratios
-        processing_queue: multiprocessing.Queue,
+    ALGORITHM_FIXED_SCALING_RATIOS = {
-        processed_frames: ListProxy,
+        "waifu2x": [1, 2],
-        pause: Synchronized,
+        "srmd": [2, 3, 4],
-    ) -> None:
+        "realsr": [4],
-        multiprocessing.Process.__init__(self)
+        "realcugan": [1, 2, 3, 4],
-        self.running = False
+    }
-        self.processing_queue = processing_queue
-        self.processed_frames = processed_frames
-        self.pause = pause
 
-        signal.signal(signal.SIGTERM, self._stop)
+    ALGORITHM_CLASSES = {
+        "waifu2x": Waifu2x,
+        "srmd": Srmd,
+        "realsr": Realsr,
+        "realcugan": Realcugan,
+    }
 
-    def run(self) -> None:
+    processor_objects = {}
-        self.running = True
+
-        logger.opt(colors=True).info(
+    @staticmethod
-            f"Upscaler process <blue>{self.name}</blue> initiating"
+    def _get_scaling_tasks(
+        input_width: int,
+        input_height: int,
+        output_width: int,
+        output_height: int,
+        algorithm: str,
+    ) -> list:
+        """
+        Get the required tasks for upscaling the image until it is larger than
+        or equal to the desired output dimensions. For example, SRMD only supports
+        2x, 3x, and 4x, so upsclaing an image from 320x240 to 3840x2160 will
+        require the SRMD to run 3x then 4x. In this case, this function will
+        return [3, 4].
+
+        :param input_width int: input image width
+        :param input_height int: input image height
+        :param output_width int: desired output image width
+        :param output_height int: desired output image size
+        :param algorithm str: upsclaing algorithm
+        :rtype list: the list of upsclaing tasks required
+        """
+        # calculate required minimum scale ratio
+        output_scale = max(output_width / input_width, output_height / input_height)
+
+        # select the optimal algorithm scaling ratio to use
+        supported_scaling_ratios = sorted(
+            Upscaler.ALGORITHM_FIXED_SCALING_RATIOS[algorithm]
         )
-        processor_objects = {}
+
-        while self.running:
+        remaining_scaling_ratio = math.ceil(output_scale)
+
+        # if the scaling ratio is 1.0
+        # apply the smallest scaling ratio available
+        if remaining_scaling_ratio == 1:
+            return [supported_scaling_ratios[0]]
+
+        scaling_jobs = []
+        while remaining_scaling_ratio > 1:
+            for ratio in supported_scaling_ratios:
+                if ratio >= remaining_scaling_ratio:
+                    scaling_jobs.append(ratio)
+                    remaining_scaling_ratio /= ratio
+                    break
+
+            else:
+                found = False
+                for i in supported_scaling_ratios:
+                    for j in supported_scaling_ratios:
+                        if i * j >= remaining_scaling_ratio:
+                            scaling_jobs.extend([i, j])
+                            remaining_scaling_ratio /= i * j
+                            found = True
+                            break
+                    if found is True:
+                        break
+
+                if found is False:
+                    scaling_jobs.append(supported_scaling_ratios[-1])
+                    remaining_scaling_ratio /= supported_scaling_ratios[-1]
+        return scaling_jobs
+
+    def upscale_image(
+        self,
+        image: Image.Image,
+        output_width: int,
+        output_height: int,
+        algorithm: str,
+        noise: int,
+    ) -> Image.Image:
+        """
+        upscale an image
+
+        :param image Image.Image: the image to upscale
+        :param output_width int: the desired output width
+        :param output_height int: the desired output height
+        :param algorithm str: the algorithm to use
+        :param noise int: the noise level (available only for some algorithms)
+        :rtype Image.Image: the upscaled image
+        """
+        width, height = image.size
+
+        for task in self._get_scaling_tasks(
+            width, height, output_width, output_height, algorithm
+        ):
+
+            # select a processor object with the required settings
+            # create a new object if none are available
+            processor_object = self.processor_objects.get((algorithm, task))
+            if processor_object is None:
+                processor_object = self.ALGORITHM_CLASSES[algorithm](
+                    noise=noise, scale=task
+                )
+                self.processor_objects[(algorithm, task)] = processor_object
+
+            # process the image with the selected algorithm
+            image = processor_object.process(image)
+
+        # downscale the image to the desired output size and
+        # save the image to disk
+        return image.resize((output_width, output_height), Image.Resampling.LANCZOS)
+
+
+class UpscalerProcessor(Processor, Upscaler):
+    def process(self) -> None:
+
+        task = self.tasks_queue.get()
+        while task is not None:
+
             try:
-                # pause if pause flag is set
+
-                if self.pause.value is True:
+                if self.pause_flag.value is True:
                     time.sleep(0.1)
                     continue
 
-                try:
+                # unpack the task's values
-                    # get new job from queue
+                (
-                    (
+                    frame_index,
-                        frame_index,
+                    previous_frame,
-                        (image0, image1),
+                    current_frame,
-                        (
+                    (output_width, output_height, algorithm, noise, threshold),
-                            output_width,
+                ) = task
-                            output_height,
-                            noise,
-                            difference_threshold,
-                            algorithm,
-                        ),
-                    ) = self.processing_queue.get(False)
-
-                    # destructure settings
-                except queue.Empty:
-                    time.sleep(0.1)
-                    continue
 
+                # calculate the %diff between the current frame and the previous frame
                 difference_ratio = 0
-                if image0 is not None:
+                if previous_frame is not None:
-                    difference = ImageChops.difference(image0, image1)
+                    difference_ratio = self.get_image_diff(
-                    difference_stat = ImageStat.Stat(difference)
+                        previous_frame, current_frame
-                    difference_ratio = (
-                        sum(difference_stat.mean)
-                        / (len(difference_stat.mean) * 255)
-                        * 100
                     )
 
-                # if the difference is lower than threshold
+                # if the difference is lower than threshold, skip this frame
-                # skip this frame
+                if difference_ratio < threshold:
-                if difference_ratio < difference_threshold:
-
-                    # make sure the previous frame has been processed
-                    if frame_index > 0:
-                        while self.processed_frames[frame_index - 1] is None:
-                            time.sleep(0.1)
 
                     # make the current image the same as the previous result
-                    self.processed_frames[frame_index] = self.processed_frames[
+                    self.processed_frames[frame_index] = True
-                        frame_index - 1
-                    ]
 
                 # if the difference is greater than threshold
                 # process this frame
                 else:
-                    width, height = image1.size
+                    self.processed_frames[frame_index] = self.upscale_image(
-
+                        current_frame, output_width, output_height, algorithm, noise
-                    # calculate required minimum scale ratio
-                    output_scale = max(output_width / width, output_height / height)
-
-                    # select the optimal algorithm scaling ratio to use
-                    supported_scaling_ratios = sorted(
-                        ALGORITHM_FIXED_SCALING_RATIOS[algorithm]
                     )
 
-                    remaining_scaling_ratio = math.ceil(output_scale)
+                task = self.tasks_queue.get()
-                    scaling_jobs = []
 
-                    # if the scaling ratio is 1.0
+            except KeyboardInterrupt:
-                    # apply the smallest scaling ratio available
-                    if remaining_scaling_ratio == 1:
-                        scaling_jobs.append(supported_scaling_ratios[0])
-                    else:
-                        while remaining_scaling_ratio > 1:
-                            for ratio in supported_scaling_ratios:
-                                if ratio >= remaining_scaling_ratio:
-                                    scaling_jobs.append(ratio)
-                                    remaining_scaling_ratio /= ratio
-                                    break
-
-                            else:
-                                found = False
-                                for i in supported_scaling_ratios:
-                                    for j in supported_scaling_ratios:
-                                        if i * j >= remaining_scaling_ratio:
-                                            scaling_jobs.extend([i, j])
-                                            remaining_scaling_ratio /= i * j
-                                            found = True
-                                            break
-                                    if found is True:
-                                        break
-
-                                if found is False:
-                                    scaling_jobs.append(supported_scaling_ratios[-1])
-                                    remaining_scaling_ratio /= supported_scaling_ratios[
-                                        -1
-                                    ]
-
-                    for job in scaling_jobs:
-
-                        # select a processor object with the required settings
-                        # create a new object if none are available
-                        processor_object = processor_objects.get((algorithm, job))
-                        if processor_object is None:
-                            processor_object = ALGORITHM_CLASSES[algorithm](
-                                noise=noise, scale=job
-                            )
-                            processor_objects[(algorithm, job)] = processor_object
-
-                        # process the image with the selected algorithm
-                        image1 = processor_object.process(image1)
-
-                    # downscale the image to the desired output size and
-                    # save the image to disk
-                    image1 = image1.resize((output_width, output_height), Image.LANCZOS)
-                    self.processed_frames[frame_index] = image1
-
-            # send exceptions into the client connection pipe
-            except (SystemExit, KeyboardInterrupt):
                 break
-
-            except Exception as e:
-                logger.exception(e)
-                break
-
-        logger.opt(colors=True).info(
-            f"Upscaler process <blue>{self.name}</blue> terminating"
-        )
-        return super().run()
-
-    def _stop(self, _signal_number, _frame) -> None:
-        self.running = False

video2x/video2x.py

@@ -27,7 +27,7 @@ __      __  _       _                  ___   __   __
 Name: Video2X
 Creator: K4YT3X
 Date Created: February 24, 2018
-Last Modified: March 21, 2022
+Last Modified: August 28, 2022
 
 Editor: BrianPetkovsek
 Last Modified: June 17, 2019
@@ -39,20 +39,18 @@ Editor: 28598519a
 Last Modified: March 23, 2020
 """
 
-import argparse
 import ctypes
 import math
-import multiprocessing
-import os
-import pathlib
 import signal
 import sys
 import time
+from enum import Enum
+from multiprocessing import Manager, Pool, Queue, Value
+from pathlib import Path
 
-import cv2
 import ffmpeg
+from cv2 import cv2
 from loguru import logger
-from rich import print
 from rich.console import Console
 from rich.file_proxy import FileProxy
 from rich.progress import (
@@ -65,43 +63,23 @@ from rich.progress import (
 )
 from rich.text import Text
 
+from video2x.processor import Processor
+
 from . import __version__
-from .decoder import VideoDecoder
+from .decoder import VideoDecoder, VideoDecoderThread
 from .encoder import VideoEncoder
 from .interpolator import Interpolator
-from .upscaler import Upscaler
+from .upscaler import UpscalerProcessor
 
 # for desktop environments only
 # if pynput can be loaded, enable global pause hotkey support
 try:
-    import pynput
+    from pynput.keyboard import HotKey, Listener
 except ImportError:
     ENABLE_HOTKEY = False
 else:
     ENABLE_HOTKEY = True
 
-LEGAL_INFO = """Video2X\t\t{}
-Author:\t\tK4YT3X
-License:\tGNU AGPL v3
-Github Page:\thttps://github.com/k4yt3x/video2x
-Contact:\ti@k4yt3x.com""".format(
-    __version__
-)
-
-# algorithms available for upscaling tasks
-UPSCALING_ALGORITHMS = [
-    "waifu2x",
-    "srmd",
-    "realsr",
-    "realcugan",
-]
-
-# algorithms available for frame interpolation tasks
-INTERPOLATION_ALGORITHMS = ["rife"]
-
-# progress bar labels for different modes
-MODE_LABELS = {"upscale": "Upscaling", "interpolate": "Interpolating"}
-
 # format string for Loguru loggers
 LOGURU_FORMAT = (
     "<green>{time:HH:mm:ss.SSSSSS!UTC}</green> | "
@@ -121,6 +99,11 @@ class ProcessingSpeedColumn(ProgressColumn):
         )
 
 
+class ProcessingMode(Enum):
+    UPSCALE = {"label": "Upscaling", "processor": UpscalerProcessor}
+    INTERPOLATE = {"label": "Interpolating", "processor": Interpolator}
+
+
 class Video2X:
     """
     Video2X class
@@ -133,11 +116,12 @@ class Video2X:
     def __init__(self) -> None:
         self.version = __version__
 
-    def _get_video_info(self, path: pathlib.Path) -> tuple:
+    @staticmethod
+    def _get_video_info(path: Path) -> tuple:
         """
         get video file information with FFmpeg
 
-        :param path pathlib.Path: video file path
+        :param path Path: video file path
         :raises RuntimeError: raised when video stream isn't found
         """
         # probe video file info
@@ -161,34 +145,17 @@ class Video2X:
 
         return video_info["width"], video_info["height"], total_frames, frame_rate
 
-    def _toggle_pause(self, _signal_number: int = -1, _frame=None):
-        # print console messages and update the progress bar's status
-        if self.pause.value is False:
-            self.progress.update(self.task, description=self.description + " (paused)")
-            self.progress.stop_task(self.task)
-            logger.warning("Processing paused, press Ctrl+Alt+V again to resume")
-
-        elif self.pause.value is True:
-            self.progress.update(self.task, description=self.description)
-            logger.warning("Resuming processing")
-            self.progress.start_task(self.task)
-
-        # invert the value of the pause flag
-        with self.pause.get_lock():
-            self.pause.value = not self.pause.value
-
     def _run(
         self,
-        input_path: pathlib.Path,
+        input_path: Path,
         width: int,
         height: int,
         total_frames: int,
         frame_rate: float,
-        output_path: pathlib.Path,
+        output_path: Path,
         output_width: int,
         output_height: int,
-        Processor: object,
+        mode: ProcessingMode,
-        mode: str,
         processes: int,
         processing_settings: tuple,
     ) -> None:
@@ -208,51 +175,40 @@ class Video2X:
         logger.remove()
         logger.add(sys.stderr, colorize=True, format=LOGURU_FORMAT)
 
-        # initialize values
+        # TODO: add docs
-        self.processor_processes = []
+        tasks_queue = Queue(maxsize=processes * 10)
-        self.processing_queue = multiprocessing.Queue(maxsize=processes * 10)
+        processed_frames = Manager().dict()
-        processed_frames = multiprocessing.Manager().list([None] * total_frames)
+        pause_flag = Value(ctypes.c_bool, False)
-        self.processed = multiprocessing.Value("I", 0)
-        self.pause = multiprocessing.Value(ctypes.c_bool, False)
 
         # set up and start decoder thread
         logger.info("Starting video decoder")
-        self.decoder = VideoDecoder(
+        decoder = VideoDecoder(
             input_path,
             width,
             height,
             frame_rate,
-            self.processing_queue,
-            processing_settings,
-            self.pause,
         )
-        self.decoder.start()
+        decoder_thread = VideoDecoderThread(tasks_queue, decoder, processing_settings)
+        decoder_thread.start()
 
         # set up and start encoder thread
         logger.info("Starting video encoder")
-        self.encoder = VideoEncoder(
+        encoder = VideoEncoder(
             input_path,
             frame_rate * 2 if mode == "interpolate" else frame_rate,
             output_path,
             output_width,
             output_height,
-            total_frames,
-            processed_frames,
-            self.processed,
-            self.pause,
         )
-        self.encoder.start()
 
-        # create processor processes
+        # create a pool of processor processes to process the queue
-        for process_name in range(processes):
+        processor: Processor = mode.value["processor"](
-            process = Processor(self.processing_queue, processed_frames, self.pause)
+            tasks_queue, processed_frames, pause_flag
-            process.name = str(process_name)
+        )
-            process.daemon = True
+        processor_pool = Pool(processes, processor.process)
-            process.start()
-            self.processor_processes.append(process)
 
         # create progress bar
-        self.progress = Progress(
+        progress = Progress(
             "[progress.description]{task.description}",
             BarColumn(complete_style="blue", finished_style="green"),
             "[progress.percentage]{task.percentage:>3.0f}%",
@@ -265,23 +221,42 @@ class Video2X:
             speed_estimate_period=300.0,
             disable=True,
         )
+        task = progress.add_task(f"[cyan]{mode.value['label']}", total=total_frames)
 
-        self.description = f"[cyan]{MODE_LABELS.get(mode, 'Unknown')}"
+        def _toggle_pause(_signal_number: int = -1, _frame=None):
-        self.task = self.progress.add_task(self.description, total=total_frames)
+
+            # allow the closure to modify external immutable flag
+            nonlocal pause_flag
+
+            # print console messages and update the progress bar's status
+            if pause_flag.value is False:
+                progress.update(
+                    task, description=f"[cyan]{mode.value['label']} (paused)"
+                )
+                progress.stop_task(task)
+                logger.warning("Processing paused, press Ctrl+Alt+V again to resume")
+
+            # the lock is already acquired
+            elif pause_flag.value is True:
+                progress.update(task, description=f"[cyan]{mode.value['label']}")
+                logger.warning("Resuming processing")
+                progress.start_task(task)
+
+            # invert the flag
+            with pause_flag.get_lock():
+                pause_flag.value = not pause_flag.value
 
         # allow sending SIGUSR1 to pause/resume processing
-        signal.signal(signal.SIGUSR1, self._toggle_pause)
+        signal.signal(signal.SIGUSR1, _toggle_pause)
 
         # enable global pause hotkey if it's supported
         if ENABLE_HOTKEY is True:
 
             # create global pause hotkey
-            pause_hotkey = pynput.keyboard.HotKey(
+            pause_hotkey = HotKey(HotKey.parse("<ctrl>+<alt>+v"), _toggle_pause)
-                pynput.keyboard.HotKey.parse("<ctrl>+<alt>+v"), self._toggle_pause
-            )
 
             # create global keyboard input listener
-            keyboard_listener = pynput.keyboard.Listener(
+            keyboard_listener = Listener(
                 on_press=(
                     lambda key: pause_hotkey.press(keyboard_listener.canonical(key))
                 ),
@@ -294,52 +269,52 @@ class Video2X:
             keyboard_listener.start()
 
         # a temporary variable that stores the exception
-        exception = []
+        exceptions = []
 
         try:
 
-            # wait for jobs in queue to deplete
+            # let the context manager automatically stop the progress bar
-            while self.processed.value < total_frames - 1:
+            with progress:
-                time.sleep(1)
 
-                # check processor health
+                frame_index = 0
-                for process in self.processor_processes:
+                while frame_index < total_frames:
-                    if not process.is_alive():
-                        raise Exception("process died unexpectedly")
 
-                # check decoder health
+                    current_frame = processed_frames.get(frame_index)
-                if not self.decoder.is_alive() and self.decoder.exception is not None:
-                    raise Exception("decoder died unexpectedly")
 
-                # check encoder health
+                    if pause_flag.value is True or current_frame is None:
-                if not self.encoder.is_alive() and self.encoder.exception is not None:
+                        time.sleep(0.1)
-                    raise Exception("encoder died unexpectedly")
+                        continue
 
-                # show progress bar when upscale starts
+                    # show the progress bar after the processing starts
-                if self.progress.disable is True and self.processed.value > 0:
+                    # reduces speed estimation inaccuracies and print overlaps
-                    self.progress.disable = False
+                    if frame_index == 0:
-                    self.progress.start()
+                        progress.disable = False
+                        progress.start()
 
-                # update progress
+                    if current_frame is True:
-                if self.pause.value is False:
+                        encoder.write(processed_frames.get(frame_index - 1))
-                    self.progress.update(self.task, completed=self.processed.value)
 
-            self.progress.update(self.task, completed=total_frames)
+                    else:
-            self.progress.stop()
+                        encoder.write(current_frame)
-            logger.info("Processing has completed")
+
+                        if frame_index > 0:
+                            del processed_frames[frame_index - 1]
+
+                    progress.update(task, completed=frame_index + 1)
+                    frame_index += 1
 
         # if SIGTERM is received or ^C is pressed
-        # TODO: pause and continue here
+        except (SystemExit, KeyboardInterrupt) as error:
-        except (SystemExit, KeyboardInterrupt) as e:
-            self.progress.stop()
             logger.warning("Exit signal received, exiting gracefully")
             logger.warning("Press ^C again to force terminate")
-            exception.append(e)
+            exceptions.append(error)
 
-        except Exception as e:
+        except Exception as error:
-            self.progress.stop()
+            logger.exception(error)
-            logger.exception(e)
+            exceptions.append(error)
-            exception.append(e)
+
+        else:
+            logger.info("Processing has completed")
 
         finally:
 
@@ -348,31 +323,28 @@ class Video2X:
                 keyboard_listener.stop()
                 keyboard_listener.join()
 
-            # stop progress display
+            # if errors have occurred, kill the FFmpeg processes
-            self.progress.stop()
+            if len(exceptions) > 0:
+                decoder.kill()
+                encoder.kill()
 
-            # stop processor processes
+            # stop the decoder
-            logger.info("Stopping processor processes")
+            decoder_thread.stop()
-            for process in self.processor_processes:
+            decoder_thread.join()
-                process.terminate()
 
-            # wait for processes to finish
+            # clear queue and signal processors to exit
-            for process in self.processor_processes:
+            # multiprocessing.Queue has no Queue.queue.clear
-                process.join()
+            while tasks_queue.empty() is not True:
+                tasks_queue.get()
+            for _ in range(processes):
+                tasks_queue.put(None)
 
-            # stop encoder and decoder
+            # close and join the process pool
-            logger.info("Stopping decoder and encoder threads")
+            processor_pool.close()
-            self.decoder.stop()
+            processor_pool.join()
-            self.encoder.stop()
-            self.decoder.join()
-            self.encoder.join()
 
-            # mark processing queue as closed
+            # stop the encoder
-            self.processing_queue.close()
+            encoder.join()
-
-            # raise the error if there is any
-            if len(exception) > 0:
-                raise exception[0]
 
             # restore original STDOUT and STDERR
             sys.stdout = original_stdout
@@ -382,10 +354,14 @@ class Video2X:
             logger.remove()
             logger.add(sys.stderr, colorize=True, format=LOGURU_FORMAT)
 
+            # raise the first collected error
+            if len(exceptions) > 0:
+                raise exceptions[0]
+
     def upscale(
         self,
-        input_path: pathlib.Path,
+        input_path: Path,
-        output_path: pathlib.Path,
+        output_path: Path,
         output_width: int,
         output_height: int,
         noise: int,
@@ -418,22 +394,21 @@ class Video2X:
             output_path,
             output_width,
             output_height,
-            Upscaler,
+            ProcessingMode.UPSCALE,
-            "upscale",
             processes,
             (
                 output_width,
                 output_height,
+                algorithm,
                 noise,
                 threshold,
-                algorithm,
             ),
         )
 
     def interpolate(
         self,
-        input_path: pathlib.Path,
+        input_path: Path,
-        output_path: pathlib.Path,
+        output_path: Path,
         processes: int,
         threshold: float,
         algorithm: str,
@@ -455,189 +430,7 @@ class Video2X:
             output_path,
             width,
             height,
-            Interpolator,
+            ProcessingMode.INTERPOLATE,
-            "interpolate",
             processes,
             (threshold, algorithm),
-        )
+        )
-
-
-def parse_arguments() -> argparse.Namespace:
-    """
-    parse command line arguments
-
-    :rtype argparse.Namespace: command parsing results
-    """
-    parser = argparse.ArgumentParser(
-        prog="video2x",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument(
-        "--version", help="show version information and exit", action="store_true"
-    )
-    parser.add_argument(
-        "-i",
-        "--input",
-        type=pathlib.Path,
-        help="input file/directory path",
-        required=True,
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        type=pathlib.Path,
-        help="output file/directory path",
-        required=True,
-    )
-    parser.add_argument(
-        "-p", "--processes", type=int, help="number of processes to launch", default=1
-    )
-    parser.add_argument(
-        "-l",
-        "--loglevel",
-        choices=["trace", "debug", "info", "success", "warning", "error", "critical"],
-        default="info",
-    )
-
-    # upscaler arguments
-    action = parser.add_subparsers(
-        help="action to perform", dest="action", required=True
-    )
-
-    upscale = action.add_parser(
-        "upscale",
-        help="upscale a file",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        add_help=False,
-    )
-    upscale.add_argument(
-        "--help", action="help", help="show this help message and exit"
-    )
-    upscale.add_argument("-w", "--width", type=int, help="output width")
-    upscale.add_argument("-h", "--height", type=int, help="output height")
-    upscale.add_argument("-n", "--noise", type=int, help="denoise level", default=3)
-    upscale.add_argument(
-        "-a",
-        "--algorithm",
-        choices=UPSCALING_ALGORITHMS,
-        help="algorithm to use for upscaling",
-        default=UPSCALING_ALGORITHMS[0],
-    )
-    upscale.add_argument(
-        "-t",
-        "--threshold",
-        type=float,
-        help=(
-            "skip if the percent difference between two adjacent frames is below this"
-            " value; set to 0 to process all frames"
-        ),
-        default=0,
-    )
-
-    # interpolator arguments
-    interpolate = action.add_parser(
-        "interpolate",
-        help="interpolate frames for file",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        add_help=False,
-    )
-    interpolate.add_argument(
-        "--help", action="help", help="show this help message and exit"
-    )
-    interpolate.add_argument(
-        "-a",
-        "--algorithm",
-        choices=UPSCALING_ALGORITHMS,
-        help="algorithm to use for upscaling",
-        default=INTERPOLATION_ALGORITHMS[0],
-    )
-    interpolate.add_argument(
-        "-t",
-        "--threshold",
-        type=float,
-        help=(
-            "skip if the percent difference between two adjacent frames exceeds this"
-            " value; set to 100 to interpolate all frames"
-        ),
-        default=10,
-    )
-
-    return parser.parse_args()
-
-
-def main() -> int:
-    """
-    command line entrypoint for direct CLI invocation
-
-    :rtype int: 0 if completed successfully, else other int
-    """
-
-    try:
-        # display version and lawful informaition
-        if "--version" in sys.argv:
-            print(LEGAL_INFO)
-            return 0
-
-        # parse command line arguments
-        args = parse_arguments()
-
-        # check input/output file paths
-        if not args.input.exists():
-            logger.critical(f"Cannot find input file: {args.input}")
-            return 1
-        elif not args.input.is_file():
-            logger.critical("Input path is not a file")
-            return 1
-
-        # set logger level
-        if os.environ.get("LOGURU_LEVEL") is None:
-            os.environ["LOGURU_LEVEL"] = args.loglevel.upper()
-
-        # remove default handler
-        logger.remove()
-
-        # add new sink with custom handler
-        logger.add(sys.stderr, colorize=True, format=LOGURU_FORMAT)
-
-        # print package version and copyright notice
-        logger.opt(colors=True).info(f"<magenta>Video2X {__version__}</magenta>")
-        logger.opt(colors=True).info(
-            "<magenta>Copyright (C) 2018-2022 K4YT3X and contributors.</magenta>"
-        )
-
-        # initialize video2x object
-        video2x = Video2X()
-
-        if args.action == "upscale":
-            video2x.upscale(
-                args.input,
-                args.output,
-                args.width,
-                args.height,
-                args.noise,
-                args.processes,
-                args.threshold,
-                args.algorithm,
-            )
-
-        elif args.action == "interpolate":
-            video2x.interpolate(
-                args.input,
-                args.output,
-                args.processes,
-                args.threshold,
-                args.algorithm,
-            )
-
-    # don't print the traceback for manual terminations
-    except KeyboardInterrupt:
-        return 2
-
-    except Exception as e:
-        logger.exception(e)
-        return 1
-
-    # if no exceptions were produced
-    else:
-        logger.success("Processing completed successfully")
-        return 0