changed decoder from a thread into an iterator

This commit is contained in:
k4yt3x 2022-04-27 22:27:48 +00:00
parent 3f457907b6
commit f3eaa47ec6

View File

@ -19,22 +19,15 @@ 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
import ffmpeg
from loguru import logger
from PIL import Image
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__(
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"]
@ -102,75 +97,27 @@ 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
)
# source depleted (decoding finished)
# after the last frame has been decoded
# read will return nothing
if len(buffer) == 0:
self.stop()
continue
)
> 0
):
# convert raw bytes into image object
image = Image.frombytes(
frame = Image.frombytes(
"RGB", (self.input_width, self.input_height), buffer
)
# 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
# return this frame
yield frame
previous_image = image
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")
def __del__(self):
# flush the remaining data in STDOUT and STDERR
self.decoder.stdout.flush()
@ -190,9 +137,3 @@ class VideoDecoder(threading.Thread):
# wait for PIPE printer to exit
self.pipe_printer.stop()
self.pipe_printer.join()
logger.info("Decoder thread exiting")
return super().run()
def stop(self) -> None:
self.running = False