mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-01 10:29:09 +00:00
changed decoder from a thread into an iterator
This commit is contained in:
parent
3f457907b6
commit
f3eaa47ec6
@ -19,22 +19,15 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|||||||
Name: Video Decoder
|
Name: Video Decoder
|
||||||
Author: K4YT3X
|
Author: K4YT3X
|
||||||
Date Created: June 17, 2021
|
Date Created: June 17, 2021
|
||||||
Last Modified: March 21, 2022
|
Last Modified: April 9, 2022
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import contextlib
|
|
||||||
import multiprocessing
|
|
||||||
import os
|
import os
|
||||||
import pathlib
|
import pathlib
|
||||||
import queue
|
|
||||||
import signal
|
import signal
|
||||||
import subprocess
|
import subprocess
|
||||||
import threading
|
|
||||||
import time
|
|
||||||
from multiprocessing.sharedctypes import Synchronized
|
|
||||||
|
|
||||||
import ffmpeg
|
import ffmpeg
|
||||||
from loguru import logger
|
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
from .pipe_printer import PipePrinter
|
from .pipe_printer import PipePrinter
|
||||||
@ -51,32 +44,34 @@ 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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
input_path: pathlib.Path,
|
input_path: pathlib.Path,
|
||||||
input_width: int,
|
input_width: int,
|
||||||
input_height: int,
|
input_height: int,
|
||||||
frame_rate: float,
|
frame_rate: float,
|
||||||
processing_queue: multiprocessing.Queue,
|
pil_ignore_max_image_pixels: bool = True,
|
||||||
processing_settings: tuple,
|
|
||||||
pause: Synchronized,
|
|
||||||
ignore_max_image_pixels=True,
|
|
||||||
) -> None:
|
) -> None:
|
||||||
threading.Thread.__init__(self)
|
|
||||||
self.running = False
|
|
||||||
self.input_path = input_path
|
self.input_path = input_path
|
||||||
self.input_width = input_width
|
self.input_width = input_width
|
||||||
self.input_height = input_height
|
self.input_height = input_height
|
||||||
self.processing_queue = processing_queue
|
|
||||||
self.processing_settings = processing_settings
|
|
||||||
self.pause = pause
|
|
||||||
|
|
||||||
# this disables the "possible DDoS" warning
|
# this disables the "possible DDoS" warning
|
||||||
if ignore_max_image_pixels:
|
if pil_ignore_max_image_pixels is True:
|
||||||
Image.MAX_IMAGE_PIXELS = None
|
Image.MAX_IMAGE_PIXELS = None
|
||||||
|
|
||||||
self.exception = None
|
|
||||||
self.decoder = subprocess.Popen(
|
self.decoder = subprocess.Popen(
|
||||||
ffmpeg.compile(
|
ffmpeg.compile(
|
||||||
ffmpeg.input(input_path, r=frame_rate)["v"]
|
ffmpeg.input(input_path, r=frame_rate)["v"]
|
||||||
@ -102,75 +97,27 @@ class VideoDecoder(threading.Thread):
|
|||||||
self.pipe_printer = PipePrinter(self.decoder.stderr)
|
self.pipe_printer = PipePrinter(self.decoder.stderr)
|
||||||
self.pipe_printer.start()
|
self.pipe_printer.start()
|
||||||
|
|
||||||
def run(self) -> None:
|
def __iter__(self):
|
||||||
self.running = True
|
|
||||||
|
|
||||||
# the index of the frame
|
# continue yielding while FFmpeg continues to produce output
|
||||||
frame_index = 0
|
while (
|
||||||
|
len(
|
||||||
# create placeholder for previous frame
|
buffer := self.decoder.stdout.read(
|
||||||
# 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(
|
|
||||||
3 * self.input_width * self.input_height
|
3 * self.input_width * self.input_height
|
||||||
)
|
)
|
||||||
|
)
|
||||||
# source depleted (decoding finished)
|
> 0
|
||||||
# after the last frame has been decoded
|
):
|
||||||
# read will return nothing
|
|
||||||
if len(buffer) == 0:
|
|
||||||
self.stop()
|
|
||||||
continue
|
|
||||||
|
|
||||||
# convert raw bytes into image object
|
# convert raw bytes into image object
|
||||||
image = Image.frombytes(
|
frame = Image.frombytes(
|
||||||
"RGB", (self.input_width, self.input_height), buffer
|
"RGB", (self.input_width, self.input_height), buffer
|
||||||
)
|
)
|
||||||
|
|
||||||
# keep checking if the running flag is set to False
|
# return this frame
|
||||||
# while waiting to put the next image into the queue
|
yield frame
|
||||||
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
|
|
||||||
|
|
||||||
previous_image = image
|
def __del__(self):
|
||||||
frame_index += 1
|
|
||||||
|
|
||||||
# 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
|
# flush the remaining data in STDOUT and STDERR
|
||||||
self.decoder.stdout.flush()
|
self.decoder.stdout.flush()
|
||||||
@ -190,9 +137,3 @@ class VideoDecoder(threading.Thread):
|
|||||||
# wait for PIPE printer to exit
|
# wait for PIPE printer to exit
|
||||||
self.pipe_printer.stop()
|
self.pipe_printer.stop()
|
||||||
self.pipe_printer.join()
|
self.pipe_printer.join()
|
||||||
|
|
||||||
logger.info("Decoder thread exiting")
|
|
||||||
return super().run()
|
|
||||||
|
|
||||||
def stop(self) -> None:
|
|
||||||
self.running = False
|
|
||||||
|
Loading…
Reference in New Issue
Block a user