mirror of
https://github.com/k4yt3x/video2x.git
synced 2026-02-10 06:44:47 +08:00
Compare commits
40 Commits
5.0.0-beta
...
5.0.0-beta
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
84b730497b | ||
|
|
db0b87597d | ||
|
|
102340e2be | ||
|
|
85437a8481 | ||
|
|
176ae90bbb | ||
|
|
44238aed35 | ||
|
|
045e643867 | ||
|
|
c92805e7bc | ||
|
|
899fe3ae2d | ||
|
|
a75c2a50ca | ||
|
|
508d6ea4d0 | ||
|
|
8976dd8199 | ||
|
|
29a55e633c | ||
|
|
f7d6dc41b3 | ||
|
|
d236131134 | ||
|
|
d669654142 | ||
|
|
4b0ab5382c | ||
|
|
737646a248 | ||
|
|
9fc0aa787e | ||
|
|
a041a60d87 | ||
|
|
020fb2dc80 | ||
|
|
9a27960bf7 | ||
|
|
862b811517 | ||
|
|
e01d24c164 | ||
|
|
0a052a3a72 | ||
|
|
f3eaa47ec6 | ||
|
|
3f457907b6 | ||
|
|
a1d750e7ca | ||
|
|
22f656b800 | ||
|
|
8eeba71ece | ||
|
|
afca10a17b | ||
|
|
f976bdc1c9 | ||
|
|
51c0c38b34 | ||
|
|
f2b2e11c41 | ||
|
|
865e3bd193 | ||
|
|
e0dc8237f5 | ||
|
|
bbc1b57445 | ||
|
|
ebbe4570d5 | ||
|
|
bcb2e97f89 | ||
|
|
ba29349e65 |
@@ -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
6
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.
|
||||
|
||||
27
README.md
27
README.md
@@ -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
|
||||
|
||||
|
||||
21
licenses/LICENSE-realcugan-ncnn-vulkan
Normal file
21
licenses/LICENSE-realcugan-ncnn-vulkan
Normal 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
165
licenses/lgpl-3.0.txt
Normal 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.
|
||||
@@ -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
BIN
tests/data/test_image.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 92 KiB |
BIN
tests/data/test_image_ref.png
Normal file
BIN
tests/data/test_image_ref.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 465 KiB |
55
tests/test_upscaler.py
Normal file
55
tests/test_upscaler.py
Normal 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
18
tests/utils.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
69
video2x/processor.py
Executable 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
|
||||
"""
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
Reference in New Issue
Block a user