40 Commits

Author SHA1 Message Date
k4yt3x
84b730497b Revert "added rich console patch for Google Colab"
This reverts commit db0b87597d.
2022-09-04 21:56:40 +00:00
k4yt3x
db0b87597d added rich console patch for Google Colab 2022-09-04 21:41:57 +00:00
k4yt3x
102340e2be updated Pillow version requirements 2022-09-04 21:32:27 +00:00
k4yt3x
85437a8481 removed the walrus operator to remain compatible with Python 3.7 2022-09-04 00:06:46 +00:00
k4yt3x
176ae90bbb removed fps_mode since it causes incompatibility issues 2022-08-28 19:23:49 +00:00
k4yt3x
44238aed35 bumped vulkan image version to 1.3-470 2022-08-28 19:18:37 +00:00
k4yt3x
045e643867 pinned opencv-python to 4.5.5.64 due to bug #676 in version 4.6.0.66 2022-08-28 19:18:17 +00:00
k4yt3x
c92805e7bc added realcugan-ncnn-vulkan's license 2022-08-28 16:02:12 +00:00
k4yt3x
899fe3ae2d added realcugan-ncnn-vulkan into licenses section 2022-08-28 16:01:23 +00:00
k4yt3x
a75c2a50ca bumped version to 5.0.0-beta6 2022-08-28 16:01:07 +00:00
K4YT3X
508d6ea4d0 Merge pull request #695 from k4yt3x/feat/decouple 2022-08-28 11:51:51 -04:00
K4YT3X
8976dd8199 Merge branch 'master' into feat/decouple 2022-08-28 11:17:39 -04:00
k4yt3x
29a55e633c fixed the issue where the encoder cannot exit 2022-08-28 15:03:16 +00:00
k4yt3x
f7d6dc41b3 changed FFmpeg vsync to fps_mode 2022-08-28 03:44:27 +00:00
K4YT3X
d236131134 Merge pull request #644 from plambe/master
Fixed interpolation algorithms list
2022-05-08 17:17:12 -04:00
plambeto
d669654142 Fixed interpolation algorithms list 2022-05-08 21:16:09 +03:00
k4yt3x
4b0ab5382c added kill method to encoder 2022-05-01 08:46:33 +00:00
k4yt3x
737646a248 made processor objects a class var 2022-05-01 08:45:56 +00:00
k4yt3x
9fc0aa787e moved main to __main__; changed _run to use the new encoder and processor 2022-05-01 08:44:33 +00:00
k4yt3x
a041a60d87 added processor skeleton class 2022-05-01 08:43:13 +00:00
k4yt3x
020fb2dc80 updated value unpacking 2022-05-01 08:42:13 +00:00
k4yt3x
9a27960bf7 added decoder thread 2022-05-01 08:34:31 +00:00
k4yt3x
862b811517 moved main to __main__ 2022-05-01 08:34:00 +00:00
k4yt3x
e01d24c164 broken the upscaler into descrete parts 2022-04-28 14:33:33 +00:00
k4yt3x
0a052a3a72 simplified the encoder 2022-04-27 22:28:14 +00:00
k4yt3x
f3eaa47ec6 changed decoder from a thread into an iterator 2022-04-27 22:27:48 +00:00
K4YT3X
3f457907b6 Merge pull request #634 from snixon/checking_output
Adding check for valid output path
2022-04-13 23:33:52 -04:00
K4YT3X
a1d750e7ca combined output dir check with the previous section 2022-04-14 03:33:22 +00:00
Steve Nixon
22f656b800 Adding check for valid output path 2022-04-11 22:47:50 -07:00
k4yt3x
8eeba71ece added configs for flake8 and isort 2022-04-09 18:43:36 +00:00
k4yt3x
afca10a17b fixed #627 env vars lost in Popen processes 2022-04-09 18:40:43 +00:00
K4YT3X
f976bdc1c9 Merge pull request #629 from snixon/ffmpeg_subs_crash
ffmpeg can crash if source title has subtitles
2022-04-09 05:16:27 -04:00
Steve Nixon
51c0c38b34 ffmpeg can crash when a title has subs 2022-04-05 23:26:31 -07:00
k4yt3x
f2b2e11c41 bumped RealCUGAN to 1.0.2 2022-04-06 03:54:14 +00:00
k4yt3x
865e3bd193 fixed Flake8 issues 2022-04-06 03:39:32 +00:00
k4yt3x
e0dc8237f5 pinned pdm-pep517 version to >=0.12 2022-04-02 18:04:47 +00:00
k4yt3x
bbc1b57445 changed license field to license-expression (PEP 639) 2022-04-02 05:36:22 +00:00
k4yt3x
ebbe4570d5 check bools explicitly 2022-04-02 05:29:54 +00:00
k4yt3x
bcb2e97f89 fixed console script error 2022-04-01 07:52:28 +00:00
k4yt3x
ba29349e65 added licensing info for pynput 2022-04-01 07:21:18 +00:00
19 changed files with 955 additions and 683 deletions

View File

@@ -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"

6
NOTICE
View File

@@ -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.

View File

@@ -82,18 +82,20 @@ Copyright (c) 2018-2022 K4YT3X and contributors.
This project includes or depends on these following projects:
| Project | License |
| ------------------------------------------------------------------- | --------------- |
| [FFmpeg](https://www.ffmpeg.org/) | LGPLv2.1, GPLv2 |
| [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-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 |
| [rife-ncnn-vulkan](https://github.com/nihui/rife-ncnn-vulkan) | MIT License |
| [ffmpeg-python](https://github.com/kkroening/ffmpeg-python) | Apache-2.0 |
| [Loguru](https://github.com/Delgan/loguru) | MIT License |
| [opencv-python](https://github.com/opencv/opencv-python) | MIT License |
| [Pillow](https://github.com/python-pillow/Pillow) | HPND License |
| [Rich](https://github.com/Textualize/rich) | MIT License |
| Project | License |
| ----------------------------------------------------------------------- | --------------- |
| [FFmpeg](https://www.ffmpeg.org/) | LGPLv2.1, GPLv2 |
| [waifu2x-ncnn-vulkan](https://github.com/nihui/waifu2x-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 |
| [rife-ncnn-vulkan](https://github.com/nihui/rife-ncnn-vulkan) | MIT License |
| [realcugan-ncnn-vulkan](https://github.com/nihui/realcugan-ncnn-vulkan) | MIT License |
| [ffmpeg-python](https://github.com/kkroening/ffmpeg-python) | Apache-2.0 |
| [Loguru](https://github.com/Delgan/loguru) | MIT License |
| [opencv-python](https://github.com/opencv/opencv-python) | 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

View File

@@ -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.

165
licenses/lgpl-3.0.txt Normal file
View File

@@ -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.

View File

@@ -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",
"pillow>=9.0.1",
"opencv-python==4.5.5.64",
"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",
"pynput>=1.7.6",
"realcugan-ncnn-vulkan-python>=1.0.2",
]
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"

BIN
tests/data/test_image.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 92 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 465 KiB

55
tests/test_upscaler.py Normal file
View File

@@ -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

18
tests/utils.py Normal file
View File

@@ -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

View File

@@ -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

View File

@@ -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())

View File

@@ -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
import time
from multiprocessing.sharedctypes import Synchronized
from multiprocessing import Queue
from queue import Full
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,
processing_settings: tuple,
pause: Synchronized,
ignore_max_image_pixels=True,
pil_ignore_max_image_pixels: bool = 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:
self.running = True
def __iter__(self):
# the index of the frame
frame_index = 0
# continue yielding while FFmpeg continues to produce output
# 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
# used in interpolate mode
previous_image = None
while len(buffer) > 0:
# continue running until an exception occurs
# or all frames have been decoded
while self.running:
# convert raw bytes into image object
frame = Image.frombytes(
"RGB", (self.input_width, self.input_height), buffer
)
# pause if pause flag is set
if self.pause.value is True:
time.sleep(0.1)
continue
# return this frame
yield frame
try:
buffer = self.decoder.stdout.read(
3 * self.input_width * self.input_height
)
# read the next frame
buffer = self.decoder.stdout.read(3 * self.input_width * self.input_height)
# source depleted (decoding finished)
# after the last frame has been decoded
# read will return nothing
if len(buffer) == 0:
self.stop()
continue
# automatically self-join and clean up after iterations are done
self.join()
# convert raw bytes into image object
image = Image.frombytes(
"RGB", (self.input_width, self.input_height), buffer
)
def kill(self):
self.decoder.send_signal(signal.SIGKILL)
# keep checking if the running flag is set to False
# 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)
def join(self):
# 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

View File

@@ -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?"],
self.original["a?"] if copy_audio is True else None,
self.original["s?"] if copy_subtitle is True else None,
self.original["d?"] if copy_data is True else None,
self.original["t?"] if copy_attachments is True else None,
# original["1:v?"],
original["a?"] if copy_audio is True else None,
original["s?"] if copy_subtitle is True else None,
original["d?"] if copy_data 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:
self.running = True
frame_index = 0
while self.running and frame_index < self.total_frames:
def kill(self):
self.encoder.send_signal(signal.SIGKILL)
# pause if pause flag is set
if self.pause.value is True:
time.sleep(0.1)
continue
def write(self, frame: Image.Image) -> None:
"""
write a frame into FFmpeg encoder's STDIN
try:
image = self.processed_frames[frame_index]
if image is None:
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")
:param frame Image.Image: the Image object to use for writing
"""
self.encoder.stdin.write(frame.tobytes())
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

View File

@@ -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:
while self.running is True:
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:
logger.exception(e)
except Exception as error:
logger.exception(error)
break
logger.opt(colors=True).info(

View File

@@ -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:

69
video2x/processor.py Executable file
View File

@@ -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
"""

View File

@@ -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, ImageChops, ImageStat
from PIL import Image
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
# 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,
}
from .processor import Processor
class Upscaler(multiprocessing.Process):
def __init__(
self,
processing_queue: multiprocessing.Queue,
processed_frames: ListProxy,
pause: Synchronized,
) -> None:
multiprocessing.Process.__init__(self)
self.running = False
self.processing_queue = processing_queue
self.processed_frames = processed_frames
self.pause = pause
class Upscaler:
# fixed scaling ratios supported by the algorithms
# 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],
}
signal.signal(signal.SIGTERM, self._stop)
ALGORITHM_CLASSES = {
"waifu2x": Waifu2x,
"srmd": Srmd,
"realsr": Realsr,
"realcugan": Realcugan,
}
def run(self) -> None:
self.running = True
logger.opt(colors=True).info(
f"Upscaler process <blue>{self.name}</blue> initiating"
processor_objects = {}
@staticmethod
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:
# get new job from queue
(
frame_index,
(image0, image1),
(
output_width,
output_height,
noise,
difference_threshold,
algorithm,
),
) = self.processing_queue.get(False)
# destructure settings
except queue.Empty:
time.sleep(0.1)
continue
# unpack the task's values
(
frame_index,
previous_frame,
current_frame,
(output_width, output_height, algorithm, noise, threshold),
) = task
# calculate the %diff between the current frame and the previous frame
difference_ratio = 0
if image0 is not None:
difference = ImageChops.difference(image0, image1)
difference_stat = ImageStat.Stat(difference)
difference_ratio = (
sum(difference_stat.mean)
/ (len(difference_stat.mean) * 255)
* 100
if previous_frame is not None:
difference_ratio = self.get_image_diff(
previous_frame, current_frame
)
# if the difference is lower than threshold
# skip this frame
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)
# if the difference is lower than threshold, skip this frame
if difference_ratio < threshold:
# make the current image the same as the previous result
self.processed_frames[frame_index] = self.processed_frames[
frame_index - 1
]
self.processed_frames[frame_index] = True
# if the difference is greater than threshold
# process this frame
else:
width, height = image1.size
# 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]
self.processed_frames[frame_index] = self.upscale_image(
current_frame, output_width, output_height, algorithm, noise
)
remaining_scaling_ratio = math.ceil(output_scale)
scaling_jobs = []
task = self.tasks_queue.get()
# if the scaling ratio is 1.0
# 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):
except 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

View File

@@ -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: str,
mode: ProcessingMode,
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
self.processor_processes = []
self.processing_queue = multiprocessing.Queue(maxsize=processes * 10)
processed_frames = multiprocessing.Manager().list([None] * total_frames)
self.processed = multiprocessing.Value("I", 0)
self.pause = multiprocessing.Value(ctypes.c_bool, False)
# TODO: add docs
tasks_queue = Queue(maxsize=processes * 10)
processed_frames = Manager().dict()
pause_flag = 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
for process_name in range(processes):
process = Processor(self.processing_queue, processed_frames, self.pause)
process.name = str(process_name)
process.daemon = True
process.start()
self.processor_processes.append(process)
# create a pool of processor processes to process the queue
processor: Processor = mode.value["processor"](
tasks_queue, processed_frames, pause_flag
)
processor_pool = Pool(processes, processor.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')}"
self.task = self.progress.add_task(self.description, total=total_frames)
def _toggle_pause(_signal_number: int = -1, _frame=None):
# 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(
pynput.keyboard.HotKey.parse("<ctrl>+<alt>+v"), self._toggle_pause
)
pause_hotkey = HotKey(HotKey.parse("<ctrl>+<alt>+v"), _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
while self.processed.value < total_frames - 1:
time.sleep(1)
# let the context manager automatically stop the progress bar
with progress:
# check processor health
for process in self.processor_processes:
if not process.is_alive():
raise Exception("process died unexpectedly")
frame_index = 0
while frame_index < total_frames:
# check decoder health
if not self.decoder.is_alive() and self.decoder.exception is not None:
raise Exception("decoder died unexpectedly")
current_frame = processed_frames.get(frame_index)
# check encoder health
if not self.encoder.is_alive() and self.encoder.exception is not None:
raise Exception("encoder died unexpectedly")
if pause_flag.value is True or current_frame is None:
time.sleep(0.1)
continue
# show progress bar when upscale starts
if self.progress.disable is True and self.processed.value > 0:
self.progress.disable = False
self.progress.start()
# show the progress bar after the processing starts
# reduces speed estimation inaccuracies and print overlaps
if frame_index == 0:
progress.disable = False
progress.start()
# update progress
if self.pause.value is False:
self.progress.update(self.task, completed=self.processed.value)
if current_frame is True:
encoder.write(processed_frames.get(frame_index - 1))
self.progress.update(self.task, completed=total_frames)
self.progress.stop()
logger.info("Processing has completed")
else:
encoder.write(current_frame)
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 e:
self.progress.stop()
except (SystemExit, KeyboardInterrupt) as error:
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:
self.progress.stop()
logger.exception(e)
exception.append(e)
except Exception as error:
logger.exception(error)
exceptions.append(error)
else:
logger.info("Processing has completed")
finally:
@@ -348,31 +323,28 @@ class Video2X:
keyboard_listener.stop()
keyboard_listener.join()
# stop progress display
self.progress.stop()
# if errors have occurred, kill the FFmpeg processes
if len(exceptions) > 0:
decoder.kill()
encoder.kill()
# stop processor processes
logger.info("Stopping processor processes")
for process in self.processor_processes:
process.terminate()
# stop the decoder
decoder_thread.stop()
decoder_thread.join()
# wait for processes to finish
for process in self.processor_processes:
process.join()
# clear queue and signal processors to exit
# multiprocessing.Queue has no Queue.queue.clear
while tasks_queue.empty() is not True:
tasks_queue.get()
for _ in range(processes):
tasks_queue.put(None)
# stop encoder and decoder
logger.info("Stopping decoder and encoder threads")
self.decoder.stop()
self.encoder.stop()
self.decoder.join()
self.encoder.join()
# close and join the process pool
processor_pool.close()
processor_pool.join()
# mark processing queue as closed
self.processing_queue.close()
# raise the error if there is any
if len(exception) > 0:
raise exception[0]
# stop the encoder
encoder.join()
# 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,
output_path: pathlib.Path,
input_path: 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,
"upscale",
ProcessingMode.UPSCALE,
processes,
(
output_width,
output_height,
algorithm,
noise,
threshold,
algorithm,
),
)
def interpolate(
self,
input_path: pathlib.Path,
output_path: pathlib.Path,
input_path: Path,
output_path: Path,
processes: int,
threshold: float,
algorithm: str,
@@ -455,189 +430,7 @@ class Video2X:
output_path,
width,
height,
Interpolator,
"interpolate",
ProcessingMode.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
)