mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 06:48:21 +00:00
Merge pull request #695 from k4yt3x/feat/decouple
This commit is contained in:
K4YT3X
GitHub
commit
508d6ea4d0
@ -31,4 +31,4 @@ __version__ = "5.0.0-beta5"
|
||||
# generated by the following lines
|
||||
from .interpolator import Interpolator
|
||||
from .upscaler import Upscaler
|
||||
from .video2x import Video2X, main
|
||||
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", fps_mode="cfr")
|
||||
.global_args("-hide_banner")
|
||||
.global_args("-nostats")
|
||||
.global_args("-nostdin")
|
||||
@ -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
|
||||
|
||||
# create placeholder for previous frame
|
||||
# used in interpolate mode
|
||||
previous_image = None
|
||||
|
||||
# continue running until an exception occurs
|
||||
# or all frames have been decoded
|
||||
while self.running is True:
|
||||
|
||||
# pause if pause flag is set
|
||||
if self.pause.value is True:
|
||||
time.sleep(0.1)
|
||||
continue
|
||||
|
||||
try:
|
||||
buffer = self.decoder.stdout.read(
|
||||
# continue yielding while FFmpeg continues to produce output
|
||||
while (
|
||||
len(
|
||||
buffer := self.decoder.stdout.read(
|
||||
3 * self.input_width * self.input_height
|
||||
)
|
||||
)
|
||||
> 0
|
||||
):
|
||||
|
||||
# source depleted (decoding finished)
|
||||
# after the last frame has been decoded
|
||||
# read will return nothing
|
||||
if len(buffer) == 0:
|
||||
self.stop()
|
||||
continue
|
||||
# convert raw bytes into image object
|
||||
frame = Image.frombytes(
|
||||
"RGB", (self.input_width, self.input_height), buffer
|
||||
)
|
||||
|
||||
# convert raw bytes into image object
|
||||
image = Image.frombytes(
|
||||
"RGB", (self.input_width, self.input_height), buffer
|
||||
)
|
||||
# return this frame
|
||||
yield frame
|
||||
|
||||
# keep checking if the running flag is set to False
|
||||
# while waiting to put the next image into the queue
|
||||
while self.running is True:
|
||||
with contextlib.suppress(queue.Full):
|
||||
self.processing_queue.put(
|
||||
(
|
||||
frame_index,
|
||||
(previous_image, image),
|
||||
self.processing_settings,
|
||||
),
|
||||
timeout=0.1,
|
||||
)
|
||||
break
|
||||
# automatically self-join and clean up after iterations are done
|
||||
self.join()
|
||||
|
||||
previous_image = image
|
||||
frame_index += 1
|
||||
def kill(self):
|
||||
self.decoder.send_signal(signal.SIGKILL)
|
||||
|
||||
# most likely "not enough image data"
|
||||
except ValueError as error:
|
||||
self.exception = error
|
||||
|
||||
# ignore queue closed
|
||||
if "is closed" not in str(error):
|
||||
logger.exception(error)
|
||||
break
|
||||
|
||||
# send exceptions into the client connection pipe
|
||||
except Exception as error:
|
||||
self.exception = error
|
||||
logger.exception(error)
|
||||
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,10 +86,10 @@ 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",
|
||||
scodec="copy",
|
||||
vsync="cfr",
|
||||
fps_mode="cfr",
|
||||
pix_fmt="yuv420p",
|
||||
crf=17,
|
||||
preset="veryslow",
|
||||
@ -138,49 +118,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 error:
|
||||
self.exception = error
|
||||
logger.exception(error)
|
||||
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()
|
||||
@ -191,9 +148,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,7 +58,6 @@ class Interpolator(multiprocessing.Process):
|
||||
logger.opt(colors=True).info(
|
||||
f"Interpolator process <blue>{self.name}</blue> initiating"
|
||||
)
|
||||
processor_objects = {}
|
||||
while self.running is True:
|
||||
try:
|
||||
# pause if pause flag is set
|
||||
@ -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
|
||||
|
||||
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 is True:
|
||||
|
||||
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 error:
|
||||
logger.exception(error)
|
||||
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: April 5, 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 ffmpeg
|
||||
from cv2 import cv2
|
||||
from loguru import logger
|
||||
from rich import print as rich_print
|
||||
from rich.console import Console
|
||||
from rich.file_proxy import FileProxy
|
||||
from rich.progress import (
|
||||
@ -65,41 +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 = 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"]
|
||||
|
||||
# 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> | "
|
||||
@ -119,6 +99,11 @@ class ProcessingSpeedColumn(ProgressColumn):
|
||||
)
|
||||
|
||||
|
||||
class ProcessingMode(Enum):
|
||||
UPSCALE = {"label": "Upscaling", "processor": UpscalerProcessor}
|
||||
INTERPOLATE = {"label": "Interpolating", "processor": Interpolator}
|
||||
|
||||
|
||||
class Video2X:
|
||||
"""
|
||||
Video2X class
|
||||
@ -132,11 +117,11 @@ class Video2X:
|
||||
self.version = __version__
|
||||
|
||||
@staticmethod
|
||||
def _get_video_info(path: pathlib.Path) -> tuple:
|
||||
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
|
||||
@ -160,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:
|
||||
@ -207,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}%",
|
||||
@ -264,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))
|
||||
),
|
||||
@ -293,51 +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
|
||||
except (SystemExit, KeyboardInterrupt) as error:
|
||||
self.progress.stop()
|
||||
logger.warning("Exit signal received, exiting gracefully")
|
||||
logger.warning("Press ^C again to force terminate")
|
||||
exception.append(error)
|
||||
exceptions.append(error)
|
||||
|
||||
except Exception as error:
|
||||
self.progress.stop()
|
||||
logger.exception(error)
|
||||
exception.append(error)
|
||||
exceptions.append(error)
|
||||
|
||||
else:
|
||||
logger.info("Processing has completed")
|
||||
|
||||
finally:
|
||||
|
||||
@ -346,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
|
||||
@ -380,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,
|
||||
@ -416,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,
|
||||
@ -453,192 +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=INTERPOLATION_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
|
||||
)
|
||||
Loading…
Reference in New Issue
Block a user