Skip to content

Commit

Permalink
feat(video2x): updated the variables in the Video2X class for compati…
Browse files Browse the repository at this point in the history
…bility with the GUI

Signed-off-by: k4yt3x <i@k4yt3x.com>
  • Loading branch information
k4yt3x committed Sep 16, 2023
1 parent 6d934e6 commit 2299302
Showing 1 changed file with 19 additions and 14 deletions.
33 changes: 19 additions & 14 deletions video2x/video2x.py
Expand Up @@ -39,7 +39,7 @@
from enum import Enum
from multiprocessing import Manager, Pool, Queue, Value
from pathlib import Path
from typing import Any
from typing import Any, Callable, Optional

import ffmpeg
from cv2 import cv2
Expand Down Expand Up @@ -106,8 +106,9 @@ class Video2X:
- interpolate: perform motion interpolation on a file
"""

def __init__(self) -> None:
def __init__(self, progress_callback: Optional[Callable] = None) -> None:
self.version = __version__
self.progress_callback = progress_callback

@staticmethod
def _get_video_info(path: Path) -> tuple:
Expand Down Expand Up @@ -229,7 +230,7 @@ def _run(
processor_pool = Pool(processes, processor.process)

# create progress bar
progress = Progress(
self.progress = Progress(
"[progress.description]{task.description}",
BarColumn(complete_style="blue", finished_style="green"),
"[progress.percentage]{task.percentage:>3.0f}%",
Expand All @@ -242,25 +243,27 @@ def _run(
speed_estimate_period=300.0,
disable=True,
)
task = progress.add_task(f"[cyan]{mode.value['label']}", total=total_frames)
task = self.progress.add_task(
f"[cyan]{mode.value['label']}", 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(
self.progress.update(
task, description=f"[cyan]{mode.value['label']} (paused)"
)
progress.stop_task(task)
self.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']}")
self.progress.update(task, description=f"[cyan]{mode.value['label']}")
logger.warning("Resuming processing")
progress.start_task(task)
self.progress.start_task(task)

# invert the flag
with pause_flag.get_lock():
Expand Down Expand Up @@ -292,7 +295,7 @@ def _toggle_pause(_signal_number: int = -1, _frame=None):

try:
# let the context manager automatically stop the progress bar
with progress:
with self.progress:
frame_index = 0
while frame_index < total_frames:
current_frame = processed_frames.get(frame_index)
Expand All @@ -304,8 +307,8 @@ def _toggle_pause(_signal_number: int = -1, _frame=None):
# show the progress bar after the processing starts
# reduces speed estimation inaccuracies and print overlaps
if frame_index == 0:
progress.disable = False
progress.start()
self.progress.disable = False
self.progress.start()

if current_frame is True:
encoder.write(processed_frames.get(frame_index - 1))
Expand All @@ -316,7 +319,9 @@ def _toggle_pause(_signal_number: int = -1, _frame=None):
if frame_index > 0:
del processed_frames[frame_index - 1]

progress.update(task, completed=frame_index + 1)
self.progress.update(task, completed=frame_index + 1)
if self.progress_callback is not None:
self.progress_callback(frame_index + 1, total_frames)
frame_index += 1

# if SIGTERM is received or ^C is pressed
Expand Down Expand Up @@ -389,10 +394,10 @@ def upscale(
width, height, total_frames, frame_rate = self._get_video_info(input_path)

# automatically calculate output width and height if only one is given
if output_width == 0 or output_width is None:
if output_width == 0:
output_width = output_height / height * width

elif output_height == 0 or output_height is None:
elif output_height == 0:
output_height = output_width / width * height

# sanitize output width and height to be divisible by 2
Expand Down

0 comments on commit 2299302

Please sign in to comment.