예제 #1
0
	def __init__(self, system):
		"""
		A status bar for a system.
		"""
		gtk.VBox.__init__(self)
		
		self.system = system
		
		# Create a separator between us and the rest of the window
		sep = gtk.HSeparator()
		self.pack_start(sep, fill = True, expand = False)
		
		# A container for all the status entries
		self.hbox = gtk.HBox(spacing = 10)
		self.hbox.set_border_width(2)
		self.pack_start(self.hbox, fill = True, expand = False)
		
		# Current system status
		self.status_label = gtk.Label()
		self.hbox.pack_start(self.status_label, fill = True, expand = False)
		self._add_sep()
		
		# Step counter
		self.step_count_label = gtk.Label()
		self.hbox.pack_start(self.step_count_label, fill = True, expand = False)
		self._add_sep()
		
		# Show progress of background tasks
		self.progress_monitor = ProgressMonitor(self.system,
		                                        orientation = gtk.ORIENTATION_HORIZONTAL)
		self.hbox.pack_start(self.progress_monitor, fill = True, expand = True)
예제 #2
0
class Upscaler:
    """An instance of this class is a upscaler that will
    upscale all images in the given directory.

    Raises:
        Exception -- all exceptions
        ArgumentError -- if argument is not valid
    """
    def __init__(
        self,
        input_path: pathlib.Path or list,
        output_path: pathlib.Path,
        driver_settings: dict,
        ffmpeg_settings: dict,
        gifski_settings: dict,
        driver: str = "waifu2x_caffe",
        scale_ratio: float = None,
        scale_width: int = None,
        scale_height: int = None,
        processes: int = 1,
        video2x_cache_directory: pathlib.Path = pathlib.Path(
            tempfile.gettempdir()) / "video2x",
        extracted_frame_format: str = "png",
        output_file_name_format_string:
        str = "{original_file_name}_output{extension}",
        image_output_extension: str = ".png",
        video_output_extension: str = ".mp4",
        preserve_frames: bool = False,
    ):

        # required parameters
        self.input = input_path
        self.output = output_path
        self.driver_settings = driver_settings
        self.ffmpeg_settings = ffmpeg_settings
        self.gifski_settings = gifski_settings

        # optional parameters
        self.driver = driver
        self.scale_ratio = scale_ratio
        self.scale_width = scale_width
        self.scale_height = scale_height
        self.processes = processes
        self.video2x_cache_directory = video2x_cache_directory
        self.extracted_frame_format = extracted_frame_format
        self.output_file_name_format_string = output_file_name_format_string
        self.image_output_extension = image_output_extension
        self.video_output_extension = video_output_extension
        self.preserve_frames = preserve_frames

        # other internal members and signals
        self.running = False
        self.current_processing_starting_time = time.time()
        self.total_frames_upscaled = 0
        self.total_frames = 0
        self.total_files = 0
        self.total_processed = 0
        self.scaling_jobs = []
        self.current_pass = 0
        self.current_input_file = pathlib.Path()
        self.last_frame_upscaled = pathlib.Path()

    def create_temp_directories(self):
        """create temporary directories"""

        # if cache directory unspecified, use %TEMP%\video2x
        if self.video2x_cache_directory is None:
            self.video2x_cache_directory = (
                pathlib.Path(tempfile.gettempdir()) / "video2x")

        # if specified cache path exists and isn't a directory
        if (self.video2x_cache_directory.exists()
                and not self.video2x_cache_directory.is_dir()):
            Avalon.error(
                _("Specified or default cache directory is a file/link"))
            raise FileExistsError(
                "Specified or default cache directory is a file/link")

        # if cache directory doesn't exist, try creating it
        if not self.video2x_cache_directory.exists():
            try:
                Avalon.debug_info(
                    _("Creating cache directory {}").format(
                        self.video2x_cache_directory))
                self.video2x_cache_directory.mkdir(parents=True, exist_ok=True)
            except Exception as exception:
                Avalon.error(
                    _("Unable to create {}").format(
                        self.video2x_cache_directory))
                raise exception

        # create temp directories for extracted frames and upscaled frames
        self.extracted_frames = pathlib.Path(
            tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(
            _("Extracted frames are being saved to: {}").format(
                self.extracted_frames))
        self.upscaled_frames = pathlib.Path(
            tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(
            _("Upscaled frames are being saved to: {}").format(
                self.upscaled_frames))

    def cleanup_temp_directories(self):
        """delete temp directories when done"""
        if not self.preserve_frames:
            for directory in [
                    self.extracted_frames,
                    self.upscaled_frames,
                    self.video2x_cache_directory,
            ]:
                try:
                    # avalon framework cannot be used if python is shutting down
                    # therefore, plain print is used
                    print(
                        _("Cleaning up cache directory: {}").format(directory))
                    shutil.rmtree(directory)
                except FileNotFoundError:
                    pass
                except OSError:
                    print(_("Unable to delete: {}").format(directory))
                    traceback.print_exc()

    def _check_arguments(self):
        if isinstance(self.input, list):
            if self.output.exists() and not self.output.is_dir():
                Avalon.error(_("Input and output path type mismatch"))
                Avalon.error(
                    _("Input is multiple files but output is not directory"))
                raise ArgumentError("input output path type mismatch")
            for input_path in self.input:
                if not input_path.is_file() and not input_path.is_dir():
                    Avalon.error(
                        _("Input path {} is neither a file nor a directory").
                        format(input_path))
                    raise FileNotFoundError(
                        f"{input_path} is neither file nor directory")
                with contextlib.suppress(FileNotFoundError):
                    if input_path.samefile(self.output):
                        Avalon.error(
                            _("Input directory and output directory cannot be the same"
                              ))
                        raise FileExistsError(
                            "input directory and output directory are the same"
                        )

        # if input is a file
        elif self.input.is_file():
            if self.output.is_dir():
                Avalon.error(_("Input and output path type mismatch"))
                Avalon.error(_("Input is single file but output is directory"))
                raise ArgumentError("input output path type mismatch")
            if self.output.suffix == "":
                Avalon.error(_("No suffix found in output file path"))
                Avalon.error(_("Suffix must be specified"))
                raise ArgumentError("no output file suffix specified")

        # if input is a directory
        elif self.input.is_dir():
            if self.output.is_file():
                Avalon.error(_("Input and output path type mismatch"))
                Avalon.error(
                    _("Input is directory but output is existing single file"))
                raise ArgumentError("input output path type mismatch")
            with contextlib.suppress(FileNotFoundError):
                if self.input.samefile(self.output):
                    Avalon.error(
                        _("Input directory and output directory cannot be the same"
                          ))
                    raise FileExistsError(
                        "input directory and output directory are the same")

        # if input is neither
        else:
            Avalon.error(_("Input path is neither a file nor a directory"))
            raise FileNotFoundError(
                f"{self.input} is neither file nor directory")

        # check FFmpeg settings
        ffmpeg_path = pathlib.Path(self.ffmpeg_settings["ffmpeg_path"])
        if not ((pathlib.Path(ffmpeg_path / "ffmpeg.exe").is_file()
                 and pathlib.Path(ffmpeg_path / "ffprobe.exe").is_file()) or
                (pathlib.Path(ffmpeg_path / "ffmpeg").is_file()
                 and pathlib.Path(ffmpeg_path / "ffprobe").is_file())):
            Avalon.error(
                _("FFmpeg or FFprobe cannot be found under the specified path")
            )
            Avalon.error(_("Please check the configuration file settings"))
            raise FileNotFoundError(self.ffmpeg_settings["ffmpeg_path"])

        # check if driver settings
        driver_settings = copy.deepcopy(self.driver_settings)
        driver_path = driver_settings.pop("path")

        # check if driver path exists
        if not (pathlib.Path(driver_path).is_file()
                or pathlib.Path(f"{driver_path}.exe").is_file()):
            Avalon.error(
                _("Specified driver executable directory doesn't exist"))
            Avalon.error(_("Please check the configuration file settings"))
            raise FileNotFoundError(driver_path)

        # parse driver arguments using driver's parser
        # the parser will throw AttributeError if argument doesn't satisfy constraints
        try:
            driver_arguments = []
            for key in driver_settings.keys():

                value = driver_settings[key]

                if value is None or value is False:
                    continue

                else:
                    if len(key) == 1:
                        driver_arguments.append(f"-{key}")
                    else:
                        driver_arguments.append(f"--{key}")
                    # true means key is an option
                    if value is not True:
                        driver_arguments.append(str(value))

            DriverWrapperMain = getattr(
                importlib.import_module(f"wrappers.{self.driver}"),
                "WrapperMain")
            DriverWrapperMain.parse_arguments(driver_arguments)
        except AttributeError as e:
            Avalon.error(
                _("Failed to parse driver argument: {}").format(e.args[0]))
            raise e

    def _upscale_frames(self, input_directory: pathlib.Path,
                        output_directory: pathlib.Path):
        """Upscale video frames with waifu2x-caffe

        This function upscales all the frames extracted
        by ffmpeg using the waifu2x-caffe binary.

        Args:
            input_directory (pathlib.Path): directory containing frames to upscale
            output_directory (pathlib.Path): directory which upscaled frames should be exported to

        Raises:
            UnrecognizedDriverError: raised when the given driver is not recognized
            e: re-raised exception after an exception has been captured and finished processing in this scope
        """

        # initialize waifu2x driver
        if self.driver not in AVAILABLE_DRIVERS:
            raise UnrecognizedDriverError(
                _("Unrecognized driver: {}").format(self.driver))

        # list all images in the extracted frames
        frames = [(input_directory / f) for f in input_directory.iterdir()
                  if f.is_file]

        # if we have less images than processes,
        # create only the processes necessary
        if len(frames) < self.processes:
            self.processes = len(frames)

        # create a directory for each process and append directory
        # name into a list
        process_directories = []
        for process_id in range(self.processes):
            process_directory = input_directory / str(process_id)
            process_directories.append(process_directory)

            # delete old directories and create new directories
            if process_directory.is_dir():
                shutil.rmtree(process_directory)
            process_directory.mkdir(parents=True, exist_ok=True)

        # waifu2x-converter-cpp will perform multi-threading within its own process
        if self.driver in [
                "waifu2x_converter_cpp",
                "waifu2x_ncnn_vulkan",
                "srmd_ncnn_vulkan",
                "realsr_ncnn_vulkan",
                "anime4kcpp",
        ]:
            process_directories = [input_directory]

        else:
            # evenly distribute images into each directory
            # until there is none left in the directory
            for image in frames:
                # move image
                image.rename(process_directories[0] / image.name)
                # rotate list
                process_directories = (process_directories[-1:] +
                                       process_directories[:-1])

        # create driver processes and start them
        for process_directory in process_directories:
            self.process_pool.append(
                self.driver_object.upscale(process_directory,
                                           output_directory))

        # start progress bar in a different thread
        Avalon.debug_info(_("Starting progress monitor"))
        self.progress_monitor = ProgressMonitor(self, process_directories)
        self.progress_monitor.start()

        # create the clearer and start it
        Avalon.debug_info(_("Starting upscaled image cleaner"))
        self.image_cleaner = ImageCleaner(input_directory, output_directory,
                                          len(self.process_pool))
        self.image_cleaner.start()

        # wait for all process to exit
        try:
            self._wait()
        except (Exception, KeyboardInterrupt, SystemExit) as e:
            # cleanup
            Avalon.debug_info(_("Killing progress monitor"))
            self.progress_monitor.stop()

            Avalon.debug_info(_("Killing upscaled image cleaner"))
            self.image_cleaner.stop()
            raise e

        # if the driver is waifu2x-converter-cpp
        # images need to be renamed to be recognizable for FFmpeg
        if self.driver == "waifu2x_converter_cpp":
            for image in [
                    f for f in output_directory.iterdir() if f.is_file()
            ]:
                renamed = re.sub(
                    f"_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}",
                    f".{self.extracted_frame_format}",
                    str(image.name),
                )
                (output_directory / image).rename(output_directory / renamed)

        # upscaling done, kill helper threads
        Avalon.debug_info(_("Killing progress monitor"))
        self.progress_monitor.stop()

        Avalon.debug_info(_("Killing upscaled image cleaner"))
        self.image_cleaner.stop()

    def _terminate_subprocesses(self):
        Avalon.warning(_("Terminating all processes"))
        for process in self.process_pool:
            process.terminate()

    def _wait(self):
        """wait for subprocesses in process pool to complete"""
        Avalon.debug_info(_("Main process waiting for subprocesses to exit"))

        try:
            # while process pool not empty
            while self.process_pool:

                # if stop signal received, terminate all processes
                if self.running is False:
                    raise SystemExit

                for process in self.process_pool:
                    process_status = process.poll()

                    # if process finished
                    if process_status is None:
                        continue

                    # if return code is not 0
                    elif process_status != 0:
                        Avalon.error(
                            _("Subprocess {} exited with code {}").format(
                                process.pid, process_status))
                        raise subprocess.CalledProcessError(
                            process_status, process.args)

                    else:
                        Avalon.debug_info(
                            _("Subprocess {} exited with code {}").format(
                                process.pid, process_status))
                        self.process_pool.remove(process)

                time.sleep(0.1)

        except (KeyboardInterrupt, SystemExit) as e:
            Avalon.warning(_("Stop signal received"))
            self._terminate_subprocesses()
            raise e

        except (Exception, subprocess.CalledProcessError) as e:
            Avalon.error(_("Subprocess execution ran into an error"))
            self._terminate_subprocesses()
            raise e

    def run(self):
        """Main controller for Video2X

        This function controls the flow of video conversion
        and handles all necessary functions.
        """

        # external stop signal when called in a thread
        self.running = True

        # define process pool to contain processes
        self.process_pool = []

        # load driver modules
        DriverWrapperMain = getattr(
            importlib.import_module(f"wrappers.{self.driver}"), "WrapperMain")
        self.driver_object = DriverWrapperMain(self.driver_settings)

        # load options from upscaler class into driver settings
        self.driver_object.load_configurations(self)

        # initialize FFmpeg object
        self.ffmpeg_object = Ffmpeg(
            self.ffmpeg_settings,
            extracted_frame_format=self.extracted_frame_format)

        # define processing queue
        self.processing_queue = queue.Queue()

        Avalon.info(_("Loading files into processing queue"))
        Avalon.debug_info(_("Input path(s): {}").format(self.input))

        # make output directory if the input is a list or a directory
        if isinstance(self.input, list) or self.input.is_dir():
            self.output.mkdir(parents=True, exist_ok=True)

        input_files = []

        # if input is single directory
        # put it in a list for compability with the following code
        if not isinstance(self.input, list):
            input_paths = [self.input]
        else:
            input_paths = self.input

        # flatten directories into file paths
        for input_path in input_paths:

            # if the input path is a single file
            # add the file's path object to input_files
            if input_path.is_file():
                input_files.append(input_path)

            # if the input path is a directory
            # add all files under the directory into the input_files (non-recursive)
            elif input_path.is_dir():
                input_files.extend(
                    [f for f in input_path.iterdir() if f.is_file()])

        output_paths = []

        for input_path in input_files:

            # get file type
            # try python-magic if it's available
            try:
                input_file_mime_type = magic.from_file(str(
                    input_path.absolute()),
                                                       mime=True)
                input_file_type = input_file_mime_type.split("/")[0]
                input_file_subtype = input_file_mime_type.split("/")[1]
            except Exception:
                input_file_mime_type = input_file_type = input_file_subtype = ""

            # if python-magic doesn't determine the file to be an image/video file
            # fall back to mimetypes to guess the file type based on the extension
            if input_file_type not in ["image", "video"]:
                # in case python-magic fails to detect file type
                # try guessing file mime type with mimetypes
                input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
                input_file_type = input_file_mime_type.split("/")[0]
                input_file_subtype = input_file_mime_type.split("/")[1]

            Avalon.debug_info(
                _("File MIME type: {}").format(input_file_mime_type))

            # set default output file suffixes
            # if image type is GIF, default output suffix is also .gif
            if input_file_mime_type == "image/gif":
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem, extension=".gif")

            elif input_file_type == "image":
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem,
                    extension=self.image_output_extension,
                )

            elif input_file_type == "video":
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem,
                    extension=self.video_output_extension,
                )

            # if file is none of: image, image/gif, video
            # skip to the next task
            else:
                Avalon.error(
                    _("File {} ({}) neither an image nor a video").format(
                        input_path, input_file_mime_type))
                Avalon.warning(_("Skipping this file"))
                continue

            # if there is only one input file
            # do not modify output file suffix
            if isinstance(self.input, pathlib.Path) and self.input.is_file():
                output_path = self.output

            output_path_id = 0
            while str(output_path) in output_paths:
                output_path = output_path.parent / pathlib.Path(
                    f"{output_path.stem}_{output_path_id}{output_path.suffix}")
                output_path_id += 1

            # record output path
            output_paths.append(str(output_path))

            # push file information into processing queue
            self.processing_queue.put((
                input_path.absolute(),
                output_path.absolute(),
                input_file_mime_type,
                input_file_type,
                input_file_subtype,
            ))

        # check argument sanity before running
        self._check_arguments()

        # record file count for external calls
        self.total_files = self.processing_queue.qsize()

        Avalon.info(_("Loaded files into processing queue"))
        # print all files in queue for debugging
        for job in self.processing_queue.queue:
            Avalon.debug_info(_("Input file: {}").format(job[0].absolute()))

        try:
            while not self.processing_queue.empty():

                # get new job from queue
                (
                    self.current_input_file,
                    output_path,
                    input_file_mime_type,
                    input_file_type,
                    input_file_subtype,
                ) = self.processing_queue.get()

                # get current job starting time for GUI calculations
                self.current_processing_starting_time = time.time()

                # get video information JSON using FFprobe
                Avalon.info(_("Reading file information"))
                file_info = self.ffmpeg_object.probe_file_info(
                    self.current_input_file)

                # create temporary directories for storing frames
                self.create_temp_directories()

                # start handling input
                # if input file is a static image
                if input_file_type == "image" and input_file_subtype != "gif":
                    Avalon.info(_("Starting upscaling image"))

                    # copy original file into the pre-processing directory
                    shutil.copy(
                        self.current_input_file,
                        self.extracted_frames / self.current_input_file.name,
                    )

                    width = int(file_info["streams"][0]["width"])
                    height = int(file_info["streams"][0]["height"])
                    framerate = self.total_frames = 1

                # elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
                else:
                    Avalon.info(_("Starting upscaling video/GIF"))

                    # find index of video stream
                    video_stream_index = None
                    for stream in file_info["streams"]:
                        if stream["codec_type"] == "video":
                            video_stream_index = stream["index"]
                            break

                    # exit if no video stream found
                    if video_stream_index is None:
                        Avalon.error(_("Aborting: No video stream found"))
                        raise StreamNotFoundError("no video stream found")

                    # get average frame rate of video stream
                    framerate = float(
                        Fraction(file_info["streams"][video_stream_index]
                                 ["r_frame_rate"]))
                    width = int(
                        file_info["streams"][video_stream_index]["width"])
                    height = int(
                        file_info["streams"][video_stream_index]["height"])

                    # get total number of frames
                    Avalon.info(
                        _("Getting total number of frames in the file"))

                    # if container stores total number of frames in nb_frames, fetch it directly
                    if "nb_frames" in file_info["streams"][video_stream_index]:
                        self.total_frames = int(
                            file_info["streams"][video_stream_index]
                            ["nb_frames"])

                    # otherwise call FFprobe to count the total number of frames
                    else:
                        self.total_frames = self.ffmpeg_object.get_number_of_frames(
                            self.current_input_file, video_stream_index)

                # calculate scale width/height/ratio and scaling jobs if required
                Avalon.info(_("Calculating scaling parameters"))

                # create a local copy of the global output settings
                output_scale = self.scale_ratio
                output_width = self.scale_width
                output_height = self.scale_height

                # calculate output width and height if scale ratio is specified
                if output_scale is not None:
                    output_width = int(
                        math.ceil(width * output_scale / 2.0) * 2)
                    output_height = int(
                        math.ceil(height * output_scale / 2.0) * 2)

                else:
                    # scale keeping aspect ratio is only one of width/height is given
                    if output_width == 0 or output_width is None:
                        output_width = output_height / height * width

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

                    output_width = int(math.ceil(output_width / 2.0) * 2)
                    output_height = int(math.ceil(output_height / 2.0) * 2)

                    # calculate required minimum scale ratio
                    output_scale = max(output_width / width,
                                       output_height / height)

                # if driver is one of the drivers that doesn't support arbitrary scaling ratio
                # TODO: more documentations on this block
                if self.driver in DRIVER_FIXED_SCALING_RATIOS:

                    # select the optimal driver scaling ratio to use
                    supported_scaling_ratios = sorted(
                        DRIVER_FIXED_SCALING_RATIOS[self.driver])

                    remaining_scaling_ratio = math.ceil(output_scale)
                    self.scaling_jobs = []

                    # if the scaling ratio is 1.0
                    # apply the smallest scaling ratio available
                    if remaining_scaling_ratio == 1:
                        self.scaling_jobs.append(supported_scaling_ratios[0])
                    else:
                        while remaining_scaling_ratio > 1:
                            for ratio in supported_scaling_ratios:
                                if ratio >= remaining_scaling_ratio:
                                    self.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:
                                            self.scaling_jobs.extend([i, j])
                                            remaining_scaling_ratio /= i * j
                                            found = True
                                            break
                                    if found is True:
                                        break

                                if found is False:
                                    self.scaling_jobs.append(
                                        supported_scaling_ratios[-1])
                                    remaining_scaling_ratio /= supported_scaling_ratios[
                                        -1]

                else:
                    self.scaling_jobs = [output_scale]

                # print file information
                Avalon.debug_info(_("Framerate: {}").format(framerate))
                Avalon.debug_info(_("Width: {}").format(width))
                Avalon.debug_info(_("Height: {}").format(height))
                Avalon.debug_info(
                    _("Total number of frames: {}").format(self.total_frames))
                Avalon.debug_info(_("Output width: {}").format(output_width))
                Avalon.debug_info(_("Output height: {}").format(output_height))
                Avalon.debug_info(
                    _("Required scale ratio: {}").format(output_scale))
                Avalon.debug_info(
                    _("Upscaling jobs queue: {}").format(self.scaling_jobs))

                # extract frames from video
                if input_file_mime_type == "image/gif" or input_file_type == "video":
                    self.process_pool.append(
                        (self.ffmpeg_object.extract_frames(
                            self.current_input_file, self.extracted_frames)))
                    self._wait()

                # if driver is waifu2x-caffe
                # pass pixel format output depth information
                if self.driver == "waifu2x_caffe":
                    # get a dict of all pixel formats and corresponding bit depth
                    pixel_formats = self.ffmpeg_object.get_pixel_formats()

                    # try getting pixel format's corresponding bti depth
                    try:
                        self.driver_settings["output_depth"] = pixel_formats[
                            self.ffmpeg_object.pixel_format]
                    except KeyError:
                        Avalon.error(
                            _("Unsupported pixel format: {}").format(
                                self.ffmpeg_object.pixel_format))
                        raise UnsupportedPixelError(
                            f"unsupported pixel format {self.ffmpeg_object.pixel_format}"
                        )

                # upscale images one by one using waifu2x
                Avalon.info(_("Starting to upscale extracted frames"))
                upscale_begin_time = time.time()

                self.current_pass = 1
                if self.driver == "waifu2x_caffe":
                    self.driver_object.set_scale_resolution(
                        output_width, output_height)
                else:
                    self.driver_object.set_scale_ratio(self.scaling_jobs[0])
                self._upscale_frames(self.extracted_frames,
                                     self.upscaled_frames)
                for job in self.scaling_jobs[1:]:
                    self.current_pass += 1
                    self.driver_object.set_scale_ratio(job)
                    shutil.rmtree(self.extracted_frames)
                    shutil.move(self.upscaled_frames, self.extracted_frames)
                    self.upscaled_frames.mkdir(parents=True, exist_ok=True)
                    self._upscale_frames(self.extracted_frames,
                                         self.upscaled_frames)

                Avalon.info(_("Upscaling completed"))
                Avalon.info(
                    _("Average processing speed: {} seconds per frame").format(
                        self.total_frames /
                        (time.time() - upscale_begin_time)))

                # downscale frames with Lanczos
                Avalon.info(_("Lanczos downscaling frames"))
                shutil.rmtree(self.extracted_frames)
                shutil.move(self.upscaled_frames, self.extracted_frames)
                self.upscaled_frames.mkdir(parents=True, exist_ok=True)

                for image in tqdm(
                    [
                        i for i in self.extracted_frames.iterdir()
                        if i.is_file()
                        and i.name.endswith(self.extracted_frame_format)
                    ],
                        ascii=True,
                        desc=_("Downscaling"),
                ):
                    image_object = Image.open(image)

                    # if the image dimensions are not equal to the output size
                    # resize the image using Lanczos
                    if (image_object.width, image_object.height) != (
                            output_width,
                            output_height,
                    ):
                        image_object.resize(
                            (output_width, output_height), Image.LANCZOS).save(
                                self.upscaled_frames / image.name)
                        image_object.close()

                    # if the image's dimensions are already equal to the output size
                    # move image to the finished directory
                    else:
                        image_object.close()
                        shutil.move(image, self.upscaled_frames / image.name)

                # start handling output
                # output can be either GIF or video
                if input_file_type == "image" and input_file_subtype != "gif":

                    Avalon.info(_("Exporting image"))

                    # there should be only one image in the directory
                    shutil.move(
                        [
                            f for f in self.upscaled_frames.iterdir()
                            if f.is_file()
                        ][0],
                        output_path,
                    )

                # elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
                else:

                    # if the desired output is gif file
                    if output_path.suffix.lower() == ".gif":
                        Avalon.info(
                            _("Converting extracted frames into GIF image"))
                        gifski_object = Gifski(self.gifski_settings)
                        self.process_pool.append(
                            gifski_object.make_gif(
                                self.upscaled_frames,
                                output_path,
                                framerate,
                                self.extracted_frame_format,
                                output_width,
                                output_height,
                            ))
                        self._wait()
                        Avalon.info(_("Conversion completed"))

                    # if the desired output is video
                    else:
                        # frames to video
                        Avalon.info(
                            _("Converting extracted frames into video"))
                        self.process_pool.append(
                            self.ffmpeg_object.assemble_video(
                                framerate, self.upscaled_frames))
                        # f'{scale_width}x{scale_height}'
                        self._wait()
                        Avalon.info(_("Conversion completed"))

                        try:
                            # migrate audio tracks and subtitles
                            Avalon.info(
                                _("Migrating audio, subtitles and other streams to upscaled video"
                                  ))
                            self.process_pool.append(
                                self.ffmpeg_object.migrate_streams(
                                    self.current_input_file,
                                    output_path,
                                    self.upscaled_frames,
                                ))
                            self._wait()

                        # if failed to copy streams
                        # use file with only video stream
                        except subprocess.CalledProcessError:
                            traceback.print_exc()
                            Avalon.error(_("Failed to migrate streams"))
                            Avalon.warning(
                                _("Trying to output video without additional streams"
                                  ))

                            if input_file_mime_type == "image/gif":
                                # copy will overwrite destination content if exists
                                shutil.copy(
                                    self.upscaled_frames /
                                    self.ffmpeg_object.intermediate_file_name,
                                    output_path,
                                )

                            else:
                                # construct output file path
                                output_file_name = f"{output_path.stem}{self.ffmpeg_object.intermediate_file_name.suffix}"
                                output_video_path = (output_path.parent /
                                                     output_file_name)

                                # if output file already exists
                                # create temporary directory in output folder
                                # temporary directories generated by tempfile are guaranteed to be unique
                                # and won't conflict with other files
                                if output_video_path.exists():
                                    Avalon.error(_("Output video file exists"))

                                    temporary_directory = pathlib.Path(
                                        tempfile.mkdtemp(
                                            dir=output_path.parent))
                                    output_video_path = (temporary_directory /
                                                         output_file_name)
                                    Avalon.info(
                                        _("Created temporary directory to contain file"
                                          ))

                                # move file to new destination
                                Avalon.info(
                                    _("Writing intermediate file to: {}").
                                    format(output_video_path.absolute()))
                                shutil.move(
                                    self.upscaled_frames /
                                    self.ffmpeg_object.intermediate_file_name,
                                    output_video_path,
                                )

                # increment total number of files processed
                self.cleanup_temp_directories()
                self.processing_queue.task_done()
                self.total_processed += 1

        except (Exception, KeyboardInterrupt, SystemExit) as e:
            with contextlib.suppress(ValueError, AttributeError):
                self.cleanup_temp_directories()
                self.running = False
            raise e

        # signal upscaling completion
        self.running = False
예제 #3
0
    def _upscale_frames(self, input_directory: pathlib.Path,
                        output_directory: pathlib.Path):
        """Upscale video frames with waifu2x-caffe

        This function upscales all the frames extracted
        by ffmpeg using the waifu2x-caffe binary.

        Args:
            input_directory (pathlib.Path): directory containing frames to upscale
            output_directory (pathlib.Path): directory which upscaled frames should be exported to

        Raises:
            UnrecognizedDriverError: raised when the given driver is not recognized
            e: re-raised exception after an exception has been captured and finished processing in this scope
        """

        # initialize waifu2x driver
        if self.driver not in AVAILABLE_DRIVERS:
            raise UnrecognizedDriverError(
                _("Unrecognized driver: {}").format(self.driver))

        # list all images in the extracted frames
        frames = [(input_directory / f) for f in input_directory.iterdir()
                  if f.is_file]

        # if we have less images than processes,
        # create only the processes necessary
        if len(frames) < self.processes:
            self.processes = len(frames)

        # create a directory for each process and append directory
        # name into a list
        process_directories = []
        for process_id in range(self.processes):
            process_directory = input_directory / str(process_id)
            process_directories.append(process_directory)

            # delete old directories and create new directories
            if process_directory.is_dir():
                shutil.rmtree(process_directory)
            process_directory.mkdir(parents=True, exist_ok=True)

        # waifu2x-converter-cpp will perform multi-threading within its own process
        if self.driver in [
                "waifu2x_converter_cpp",
                "waifu2x_ncnn_vulkan",
                "srmd_ncnn_vulkan",
                "realsr_ncnn_vulkan",
                "anime4kcpp",
        ]:
            process_directories = [input_directory]

        else:
            # evenly distribute images into each directory
            # until there is none left in the directory
            for image in frames:
                # move image
                image.rename(process_directories[0] / image.name)
                # rotate list
                process_directories = (process_directories[-1:] +
                                       process_directories[:-1])

        # create driver processes and start them
        for process_directory in process_directories:
            self.process_pool.append(
                self.driver_object.upscale(process_directory,
                                           output_directory))

        # start progress bar in a different thread
        Avalon.debug_info(_("Starting progress monitor"))
        self.progress_monitor = ProgressMonitor(self, process_directories)
        self.progress_monitor.start()

        # create the clearer and start it
        Avalon.debug_info(_("Starting upscaled image cleaner"))
        self.image_cleaner = ImageCleaner(input_directory, output_directory,
                                          len(self.process_pool))
        self.image_cleaner.start()

        # wait for all process to exit
        try:
            self._wait()
        except (Exception, KeyboardInterrupt, SystemExit) as e:
            # cleanup
            Avalon.debug_info(_("Killing progress monitor"))
            self.progress_monitor.stop()

            Avalon.debug_info(_("Killing upscaled image cleaner"))
            self.image_cleaner.stop()
            raise e

        # if the driver is waifu2x-converter-cpp
        # images need to be renamed to be recognizable for FFmpeg
        if self.driver == "waifu2x_converter_cpp":
            for image in [
                    f for f in output_directory.iterdir() if f.is_file()
            ]:
                renamed = re.sub(
                    f"_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}",
                    f".{self.extracted_frame_format}",
                    str(image.name),
                )
                (output_directory / image).rename(output_directory / renamed)

        # upscaling done, kill helper threads
        Avalon.debug_info(_("Killing progress monitor"))
        self.progress_monitor.stop()

        Avalon.debug_info(_("Killing upscaled image cleaner"))
        self.image_cleaner.stop()
예제 #4
0
class Upscaler:
    """ An instance of this class is a upscaler that will
    upscale all images in the given directory.

    Raises:
        Exception -- all exceptions
        ArgumentError -- if argument is not valid
    """
    def __init__(self,
                 input_path: pathlib.Path or list,
                 output_path: pathlib.Path,
                 driver_settings: dict,
                 ffmpeg_settings: dict,
                 gifski_settings: dict,
                 driver: str = 'waifu2x_caffe',
                 scale_ratio: float = None,
                 processes: int = 1,
                 video2x_cache_directory: pathlib.Path = pathlib.Path(
                     tempfile.gettempdir()) / 'video2x',
                 extracted_frame_format: str = 'png',
                 output_file_name_format_string:
                 str = '{original_file_name}_output{extension}',
                 image_output_extension: str = '.png',
                 video_output_extension: str = '.mp4',
                 preserve_frames: bool = False):

        # required parameters
        self.input = input_path
        self.output = output_path
        self.driver_settings = driver_settings
        self.ffmpeg_settings = ffmpeg_settings
        self.gifski_settings = gifski_settings

        # optional parameters
        self.driver = driver
        self.scale_ratio = scale_ratio
        self.processes = processes
        self.video2x_cache_directory = video2x_cache_directory
        self.extracted_frame_format = extracted_frame_format
        self.output_file_name_format_string = output_file_name_format_string
        self.image_output_extension = image_output_extension
        self.video_output_extension = video_output_extension
        self.preserve_frames = preserve_frames

        # other internal members and signals
        self.running = False
        self.total_frames_upscaled = 0
        self.total_frames = 0
        self.total_files = 0
        self.total_processed = 0
        self.current_input_file = pathlib.Path()
        self.last_frame_upscaled = pathlib.Path()

    def create_temp_directories(self):
        """create temporary directories
        """

        # if cache directory unspecified, use %TEMP%\video2x
        if self.video2x_cache_directory is None:
            self.video2x_cache_directory = pathlib.Path(
                tempfile.gettempdir()) / 'video2x'

        # if specified cache path exists and isn't a directory
        if self.video2x_cache_directory.exists(
        ) and not self.video2x_cache_directory.is_dir():
            Avalon.error(
                _('Specified or default cache directory is a file/link'))
            raise FileExistsError(
                'Specified or default cache directory is a file/link')

        # if cache directory doesn't exist, try creating it
        if not self.video2x_cache_directory.exists():
            try:
                Avalon.debug_info(
                    _('Creating cache directory {}').format(
                        self.video2x_cache_directory))
                self.video2x_cache_directory.mkdir(parents=True, exist_ok=True)
            except Exception as exception:
                Avalon.error(
                    _('Unable to create {}').format(
                        self.video2x_cache_directory))
                raise exception

        # create temp directories for extracted frames and upscaled frames
        self.extracted_frames = pathlib.Path(
            tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(
            _('Extracted frames are being saved to: {}').format(
                self.extracted_frames))
        self.upscaled_frames = pathlib.Path(
            tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(
            _('Upscaled frames are being saved to: {}').format(
                self.upscaled_frames))

    def cleanup_temp_directories(self):
        """delete temp directories when done
        """
        if not self.preserve_frames:
            for directory in [
                    self.extracted_frames, self.upscaled_frames,
                    self.video2x_cache_directory
            ]:
                try:
                    # avalon framework cannot be used if python is shutting down
                    # therefore, plain print is used
                    print(
                        _('Cleaning up cache directory: {}').format(directory))
                    shutil.rmtree(directory)
                except FileNotFoundError:
                    pass
                except OSError:
                    print(_('Unable to delete: {}').format(directory))
                    traceback.print_exc()

    def _check_arguments(self):
        if isinstance(self.input, list):
            if self.output.exists() and not self.output.is_dir():
                Avalon.error(_('Input and output path type mismatch'))
                Avalon.error(
                    _('Input is multiple files but output is not directory'))
                raise ArgumentError('input output path type mismatch')
            for input_path in self.input:
                if not input_path.is_file() and not input_path.is_dir():
                    Avalon.error(
                        _('Input path {} is neither a file nor a directory').
                        format(input_path))
                    raise FileNotFoundError(
                        f'{input_path} is neither file nor directory')
                with contextlib.suppress(FileNotFoundError):
                    if input_path.samefile(self.output):
                        Avalon.error(
                            _('Input directory and output directory cannot be the same'
                              ))
                        raise FileExistsError(
                            'input directory and output directory are the same'
                        )

        # if input is a file
        elif self.input.is_file():
            if self.output.is_dir():
                Avalon.error(_('Input and output path type mismatch'))
                Avalon.error(_('Input is single file but output is directory'))
                raise ArgumentError('input output path type mismatch')
            if self.output.suffix == '':
                Avalon.error(_('No suffix found in output file path'))
                Avalon.error(_('Suffix must be specified'))
                raise ArgumentError('no output file suffix specified')

        # if input is a directory
        elif self.input.is_dir():
            if self.output.is_file():
                Avalon.error(_('Input and output path type mismatch'))
                Avalon.error(
                    _('Input is directory but output is existing single file'))
                raise ArgumentError('input output path type mismatch')
            with contextlib.suppress(FileNotFoundError):
                if self.input.samefile(self.output):
                    Avalon.error(
                        _('Input directory and output directory cannot be the same'
                          ))
                    raise FileExistsError(
                        'input directory and output directory are the same')

        # if input is neither
        else:
            Avalon.error(_('Input path is neither a file nor a directory'))
            raise FileNotFoundError(
                f'{self.input} is neither file nor directory')

        # check FFmpeg settings
        ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path'])
        if not ((pathlib.Path(ffmpeg_path / 'ffmpeg.exe').is_file()
                 and pathlib.Path(ffmpeg_path / 'ffprobe.exe').is_file()) or
                (pathlib.Path(ffmpeg_path / 'ffmpeg').is_file()
                 and pathlib.Path(ffmpeg_path / 'ffprobe').is_file())):
            Avalon.error(
                _('FFmpeg or FFprobe cannot be found under the specified path')
            )
            Avalon.error(_('Please check the configuration file settings'))
            raise FileNotFoundError(self.ffmpeg_settings['ffmpeg_path'])

        # check if driver settings
        driver_settings = copy.deepcopy(self.driver_settings)
        driver_path = driver_settings.pop('path')

        # check if driver path exists
        if not (pathlib.Path(driver_path).is_file()
                or pathlib.Path(f'{driver_path}.exe').is_file()):
            Avalon.error(
                _('Specified driver executable directory doesn\'t exist'))
            Avalon.error(_('Please check the configuration file settings'))
            raise FileNotFoundError(driver_path)

        # parse driver arguments using driver's parser
        # the parser will throw AttributeError if argument doesn't satisfy constraints
        try:
            driver_arguments = []
            for key in driver_settings.keys():

                value = driver_settings[key]

                if value is None or value is False:
                    continue

                else:
                    if len(key) == 1:
                        driver_arguments.append(f'-{key}')
                    else:
                        driver_arguments.append(f'--{key}')
                    # true means key is an option
                    if value is not True:
                        driver_arguments.append(str(value))

            DriverWrapperMain = getattr(
                importlib.import_module(f'wrappers.{self.driver}'),
                'WrapperMain')
            DriverWrapperMain.parse_arguments(driver_arguments)
        except AttributeError as e:
            Avalon.error(
                _('Failed to parse driver argument: {}').format(e.args[0]))
            raise e

        # waifu2x-caffe scale_ratio, scale_width and scale_height check
        if self.driver == 'waifu2x_caffe':
            if (driver_settings['scale_width'] != 0
                    and driver_settings['scale_height'] == 0
                    or driver_settings['scale_width'] == 0
                    and driver_settings['scale_height'] != 0):
                Avalon.error(
                    _('Only one of scale_width and scale_height is specified for waifu2x-caffe'
                      ))
                raise AttributeError(
                    'only one of scale_width and scale_height is specified for waifu2x-caffe'
                )

            # if scale_width and scale_height are specified, ensure scale_ratio is None
            elif self.driver_settings[
                    'scale_width'] != 0 and self.driver_settings[
                        'scale_height'] != 0:
                self.driver_settings['scale_ratio'] = None

            # if scale_width and scale_height not specified
            # ensure they are None, not 0
            else:
                self.driver_settings['scale_width'] = None
                self.driver_settings['scale_height'] = None

    def _upscale_frames(self):
        """ Upscale video frames with waifu2x-caffe

        This function upscales all the frames extracted
        by ffmpeg using the waifu2x-caffe binary.

        Arguments:
            w2 {Waifu2x Object} -- initialized waifu2x object
        """

        # initialize waifu2x driver
        if self.driver not in AVAILABLE_DRIVERS:
            raise UnrecognizedDriverError(
                _('Unrecognized driver: {}').format(self.driver))

        # list all images in the extracted frames
        frames = [(self.extracted_frames / f)
                  for f in self.extracted_frames.iterdir() if f.is_file]

        # if we have less images than processes,
        # create only the processes necessary
        if len(frames) < self.processes:
            self.processes = len(frames)

        # create a directory for each process and append directory
        # name into a list
        process_directories = []
        for process_id in range(self.processes):
            process_directory = self.extracted_frames / str(process_id)
            process_directories.append(process_directory)

            # delete old directories and create new directories
            if process_directory.is_dir():
                shutil.rmtree(process_directory)
            process_directory.mkdir(parents=True, exist_ok=True)

        # waifu2x-converter-cpp will perform multi-threading within its own process
        if self.driver in [
                'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan',
                'srmd_ncnn_vulkan', 'realsr_ncnn_vulkan', 'anime4kcpp'
        ]:
            process_directories = [self.extracted_frames]

        else:
            # evenly distribute images into each directory
            # until there is none left in the directory
            for image in frames:
                # move image
                image.rename(process_directories[0] / image.name)
                # rotate list
                process_directories = process_directories[
                    -1:] + process_directories[:-1]

        # create driver processes and start them
        for process_directory in process_directories:
            self.process_pool.append(
                self.driver_object.upscale(process_directory,
                                           self.upscaled_frames))

        # start progress bar in a different thread
        Avalon.debug_info(_('Starting progress monitor'))
        self.progress_monitor = ProgressMonitor(self, process_directories)
        self.progress_monitor.start()

        # create the clearer and start it
        Avalon.debug_info(_('Starting upscaled image cleaner'))
        self.image_cleaner = ImageCleaner(self.extracted_frames,
                                          self.upscaled_frames,
                                          len(self.process_pool))
        self.image_cleaner.start()

        # wait for all process to exit
        try:
            self._wait()
        except (Exception, KeyboardInterrupt, SystemExit) as e:
            # cleanup
            Avalon.debug_info(_('Killing progress monitor'))
            self.progress_monitor.stop()

            Avalon.debug_info(_('Killing upscaled image cleaner'))
            self.image_cleaner.stop()
            raise e

        # if the driver is waifu2x-converter-cpp
        # images need to be renamed to be recognizable for FFmpeg
        if self.driver == 'waifu2x_converter_cpp':
            for image in [
                    f for f in self.upscaled_frames.iterdir() if f.is_file()
            ]:
                renamed = re.sub(
                    f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}',
                    f'.{self.extracted_frame_format}', str(image.name))
                (self.upscaled_frames / image).rename(self.upscaled_frames /
                                                      renamed)

        # upscaling done, kill helper threads
        Avalon.debug_info(_('Killing progress monitor'))
        self.progress_monitor.stop()

        Avalon.debug_info(_('Killing upscaled image cleaner'))
        self.image_cleaner.stop()

    def _terminate_subprocesses(self):
        Avalon.warning(_('Terminating all processes'))
        for process in self.process_pool:
            process.terminate()

    def _wait(self):
        """ wait for subprocesses in process pool to complete
        """
        Avalon.debug_info(_('Main process waiting for subprocesses to exit'))

        try:
            # while process pool not empty
            while self.process_pool:

                # if stop signal received, terminate all processes
                if self.running is False:
                    raise SystemExit

                for process in self.process_pool:
                    process_status = process.poll()

                    # if process finished
                    if process_status is None:
                        continue

                    # if return code is not 0
                    elif process_status != 0:
                        Avalon.error(
                            _('Subprocess {} exited with code {}').format(
                                process.pid, process_status))
                        raise subprocess.CalledProcessError(
                            process_status, process.args)

                    else:
                        Avalon.debug_info(
                            _('Subprocess {} exited with code {}').format(
                                process.pid, process_status))
                        self.process_pool.remove(process)

                time.sleep(0.1)

        except (KeyboardInterrupt, SystemExit) as e:
            Avalon.warning(_('Stop signal received'))
            self._terminate_subprocesses()
            raise e

        except (Exception, subprocess.CalledProcessError) as e:
            Avalon.error(_('Subprocess execution ran into an error'))
            self._terminate_subprocesses()
            raise e

    def run(self):
        """ Main controller for Video2X

        This function controls the flow of video conversion
        and handles all necessary functions.
        """

        # external stop signal when called in a thread
        self.running = True

        # define process pool to contain processes
        self.process_pool = []

        # load driver modules
        DriverWrapperMain = getattr(
            importlib.import_module(f'wrappers.{self.driver}'), 'WrapperMain')
        self.driver_object = DriverWrapperMain(self.driver_settings)

        # load options from upscaler class into driver settings
        self.driver_object.load_configurations(self)

        # initialize FFmpeg object
        self.ffmpeg_object = Ffmpeg(
            self.ffmpeg_settings,
            extracted_frame_format=self.extracted_frame_format)

        # define processing queue
        self.processing_queue = queue.Queue()

        Avalon.info(_('Loading files into processing queue'))
        Avalon.debug_info(_('Input path(s): {}').format(self.input))

        # make output directory if the input is a list or a directory
        if isinstance(self.input, list) or self.input.is_dir():
            self.output.mkdir(parents=True, exist_ok=True)

        input_files = []

        # if input is single directory
        # put it in a list for compability with the following code
        if not isinstance(self.input, list):
            input_paths = [self.input]
        else:
            input_paths = self.input

        # flatten directories into file paths
        for input_path in input_paths:

            # if the input path is a single file
            # add the file's path object to input_files
            if input_path.is_file():
                input_files.append(input_path)

            # if the input path is a directory
            # add all files under the directory into the input_files (non-recursive)
            elif input_path.is_dir():
                input_files.extend(
                    [f for f in input_path.iterdir() if f.is_file()])

        output_paths = []

        for input_path in input_files:

            # get file type
            # try python-magic if it's available
            try:
                input_file_mime_type = magic.from_file(str(
                    input_path.absolute()),
                                                       mime=True)
                input_file_type = input_file_mime_type.split('/')[0]
                input_file_subtype = input_file_mime_type.split('/')[1]
            except Exception:
                input_file_type = input_file_subtype = None

            # in case python-magic fails to detect file type
            # try guessing file mime type with mimetypes
            if input_file_type not in ['image', 'video']:
                input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
                input_file_type = input_file_mime_type.split('/')[0]
                input_file_subtype = input_file_mime_type.split('/')[1]

            # set default output file suffixes
            # if image type is GIF, default output suffix is also .gif
            if input_file_mime_type == 'image/gif':
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem, extension='.gif')

            elif input_file_type == 'image':
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem,
                    extension=self.image_output_extension)

            elif input_file_type == 'video':
                output_path = self.output / self.output_file_name_format_string.format(
                    original_file_name=input_path.stem,
                    extension=self.video_output_extension)

            # if file is none of: image, image/gif, video
            # skip to the next task
            else:
                Avalon.error(
                    _('File {} ({}) neither an image nor a video').format(
                        input_path, input_file_mime_type))
                Avalon.warning(_('Skipping this file'))
                continue

            # if there is only one input file
            # do not modify output file suffix
            if isinstance(self.input, pathlib.Path) and self.input.is_file():
                output_path = self.output

            output_path_id = 0
            while str(output_path) in output_paths:
                output_path = output_path.parent / pathlib.Path(
                    f'{output_path.stem}_{output_path_id}{output_path.suffix}')
                output_path_id += 1

            # record output path
            output_paths.append(str(output_path))

            # push file information into processing queue
            self.processing_queue.put(
                (input_path.absolute(), output_path.absolute(),
                 input_file_mime_type, input_file_type, input_file_subtype))

        # check argument sanity before running
        self._check_arguments()

        # record file count for external calls
        self.total_files = self.processing_queue.qsize()

        Avalon.info(_('Loaded files into processing queue'))
        # print all files in queue for debugging
        for job in self.processing_queue.queue:
            Avalon.debug_info(_('Input file: {}').format(job[0].absolute()))

        try:
            while not self.processing_queue.empty():

                # get new job from queue
                self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype = self.processing_queue.get(
                )

                # start handling input
                # if input file is a static image
                if input_file_type == 'image' and input_file_subtype != 'gif':
                    Avalon.info(_('Starting to upscale image'))
                    self.process_pool.append(
                        self.driver_object.upscale(self.current_input_file,
                                                   output_path))
                    self._wait()
                    Avalon.info(_('Upscaling completed'))

                    # static images don't require GIF or video encoding
                    # go to the next task
                    self.processing_queue.task_done()
                    self.total_processed += 1
                    continue

                # if input file is a image/gif file or a video
                elif input_file_mime_type == 'image/gif' or input_file_type == 'video':

                    self.create_temp_directories()

                    # get video information JSON using FFprobe
                    Avalon.info(_('Reading video information'))
                    video_info = self.ffmpeg_object.probe_file_info(
                        self.current_input_file)
                    # analyze original video with FFprobe and retrieve framerate
                    # width, height = info['streams'][0]['width'], info['streams'][0]['height']

                    # find index of video stream
                    video_stream_index = None
                    for stream in video_info['streams']:
                        if stream['codec_type'] == 'video':
                            video_stream_index = stream['index']
                            break

                    # exit if no video stream found
                    if video_stream_index is None:
                        Avalon.error(_('Aborting: No video stream found'))
                        raise StreamNotFoundError('no video stream found')

                    # get average frame rate of video stream
                    framerate = float(
                        Fraction(video_info['streams'][video_stream_index]
                                 ['r_frame_rate']))
                    Avalon.info(_('Framerate: {}').format(framerate))
                    # self.ffmpeg_object.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']

                    # extract frames from video
                    self.process_pool.append(
                        (self.ffmpeg_object.extract_frames(
                            self.current_input_file, self.extracted_frames)))
                    self._wait()

                    # if driver is waifu2x-caffe
                    # pass pixel format output depth information
                    if self.driver == 'waifu2x_caffe':
                        # get a dict of all pixel formats and corresponding bit depth
                        pixel_formats = self.ffmpeg_object.get_pixel_formats()

                        # try getting pixel format's corresponding bti depth
                        try:
                            self.driver_settings[
                                'output_depth'] = pixel_formats[
                                    self.ffmpeg_object.pixel_format]
                        except KeyError:
                            Avalon.error(
                                _('Unsupported pixel format: {}').format(
                                    self.ffmpeg_object.pixel_format))
                            raise UnsupportedPixelError(
                                f'unsupported pixel format {self.ffmpeg_object.pixel_format}'
                            )

                    # width/height will be coded width/height x upscale factor
                    # original_width = video_info['streams'][video_stream_index]['width']
                    # original_height = video_info['streams'][video_stream_index]['height']
                    # scale_width = int(self.scale_ratio * original_width)
                    # scale_height = int(self.scale_ratio * original_height)

                    # upscale images one by one using waifu2x
                    Avalon.info(_('Starting to upscale extracted frames'))
                    self._upscale_frames()
                    Avalon.info(_('Upscaling completed'))

                # start handling output
                # output can be either GIF or video

                # if the desired output is gif file
                if output_path.suffix.lower() == '.gif':
                    Avalon.info(
                        _('Converting extracted frames into GIF image'))
                    gifski_object = Gifski(self.gifski_settings)
                    self.process_pool.append(
                        gifski_object.make_gif(self.upscaled_frames,
                                               output_path, framerate,
                                               self.extracted_frame_format))
                    self._wait()
                    Avalon.info(_('Conversion completed'))

                # if the desired output is video
                else:
                    # frames to video
                    Avalon.info(_('Converting extracted frames into video'))
                    self.process_pool.append(
                        self.ffmpeg_object.assemble_video(
                            framerate, self.upscaled_frames))
                    # f'{scale_width}x{scale_height}'
                    self._wait()
                    Avalon.info(_('Conversion completed'))

                    try:
                        # migrate audio tracks and subtitles
                        Avalon.info(
                            _('Migrating audio, subtitles and other streams to upscaled video'
                              ))
                        self.process_pool.append(
                            self.ffmpeg_object.migrate_streams(
                                self.current_input_file, output_path,
                                self.upscaled_frames))
                        self._wait()

                    # if failed to copy streams
                    # use file with only video stream
                    except subprocess.CalledProcessError:
                        traceback.print_exc()
                        Avalon.error(_('Failed to migrate streams'))
                        Avalon.warning(
                            _('Trying to output video without additional streams'
                              ))

                        if input_file_mime_type == 'image/gif':
                            # copy will overwrite destination content if exists
                            shutil.copy(
                                self.upscaled_frames /
                                self.ffmpeg_object.intermediate_file_name,
                                output_path)

                        else:
                            # construct output file path
                            output_file_name = f'{output_path.stem}{self.ffmpeg_object.intermediate_file_name.suffix}'
                            output_video_path = output_path.parent / output_file_name

                            # if output file already exists
                            # create temporary directory in output folder
                            # temporary directories generated by tempfile are guaranteed to be unique
                            # and won't conflict with other files
                            if output_video_path.exists():
                                Avalon.error(_('Output video file exists'))

                                temporary_directory = pathlib.Path(
                                    tempfile.mkdtemp(dir=output_path.parent))
                                output_video_path = temporary_directory / output_file_name
                                Avalon.info(
                                    _('Created temporary directory to contain file'
                                      ))

                            # move file to new destination
                            Avalon.info(
                                _('Writing intermediate file to: {}').format(
                                    output_video_path.absolute()))
                            shutil.move(
                                self.upscaled_frames /
                                self.ffmpeg_object.intermediate_file_name,
                                output_video_path)

                # increment total number of files processed
                self.cleanup_temp_directories()
                self.processing_queue.task_done()
                self.total_processed += 1

        except (Exception, KeyboardInterrupt, SystemExit) as e:
            with contextlib.suppress(ValueError, AttributeError):
                self.cleanup_temp_directories()
                self.running = False
            raise e

        # signal upscaling completion
        self.running = False
예제 #5
0
    def _upscale_frames(self):
        """ Upscale video frames with waifu2x-caffe

        This function upscales all the frames extracted
        by ffmpeg using the waifu2x-caffe binary.

        Arguments:
            w2 {Waifu2x Object} -- initialized waifu2x object
        """

        # initialize waifu2x driver
        if self.driver not in AVAILABLE_DRIVERS:
            raise UnrecognizedDriverError(
                _('Unrecognized driver: {}').format(self.driver))

        # list all images in the extracted frames
        frames = [(self.extracted_frames / f)
                  for f in self.extracted_frames.iterdir() if f.is_file]

        # if we have less images than processes,
        # create only the processes necessary
        if len(frames) < self.processes:
            self.processes = len(frames)

        # create a directory for each process and append directory
        # name into a list
        process_directories = []
        for process_id in range(self.processes):
            process_directory = self.extracted_frames / str(process_id)
            process_directories.append(process_directory)

            # delete old directories and create new directories
            if process_directory.is_dir():
                shutil.rmtree(process_directory)
            process_directory.mkdir(parents=True, exist_ok=True)

        # waifu2x-converter-cpp will perform multi-threading within its own process
        if self.driver in [
                'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan',
                'srmd_ncnn_vulkan', 'realsr_ncnn_vulkan', 'anime4kcpp'
        ]:
            process_directories = [self.extracted_frames]

        else:
            # evenly distribute images into each directory
            # until there is none left in the directory
            for image in frames:
                # move image
                image.rename(process_directories[0] / image.name)
                # rotate list
                process_directories = process_directories[
                    -1:] + process_directories[:-1]

        # create driver processes and start them
        for process_directory in process_directories:
            self.process_pool.append(
                self.driver_object.upscale(process_directory,
                                           self.upscaled_frames))

        # start progress bar in a different thread
        Avalon.debug_info(_('Starting progress monitor'))
        self.progress_monitor = ProgressMonitor(self, process_directories)
        self.progress_monitor.start()

        # create the clearer and start it
        Avalon.debug_info(_('Starting upscaled image cleaner'))
        self.image_cleaner = ImageCleaner(self.extracted_frames,
                                          self.upscaled_frames,
                                          len(self.process_pool))
        self.image_cleaner.start()

        # wait for all process to exit
        try:
            self._wait()
        except (Exception, KeyboardInterrupt, SystemExit) as e:
            # cleanup
            Avalon.debug_info(_('Killing progress monitor'))
            self.progress_monitor.stop()

            Avalon.debug_info(_('Killing upscaled image cleaner'))
            self.image_cleaner.stop()
            raise e

        # if the driver is waifu2x-converter-cpp
        # images need to be renamed to be recognizable for FFmpeg
        if self.driver == 'waifu2x_converter_cpp':
            for image in [
                    f for f in self.upscaled_frames.iterdir() if f.is_file()
            ]:
                renamed = re.sub(
                    f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}',
                    f'.{self.extracted_frame_format}', str(image.name))
                (self.upscaled_frames / image).rename(self.upscaled_frames /
                                                      renamed)

        # upscaling done, kill helper threads
        Avalon.debug_info(_('Killing progress monitor'))
        self.progress_monitor.stop()

        Avalon.debug_info(_('Killing upscaled image cleaner'))
        self.image_cleaner.stop()
예제 #6
0
def load():
    if conf_file:
        file = open(conf_file)
        conf = yaml.load(file)
        try:
            default_password = conf['default_passwd']
            disk_thresold = conf['disk_threshold']
            disk_monitor_interval = conf['monitors']['disk_monitor']
            progress_monitor_interval = conf['monitors']['progress_monitor']
            gateway_monitor_interval = conf['monitors']['gateway_monitor']
            db_conn_monitor_interval = conf['monitors']['db_conn_monitor']
            smtp_server = conf['email']['server']
            smtp_port = conf['email']['port']
            smtp_username = conf['email']['username']
            smtp_password = conf['email']['passwd']
            mailto_list = conf['email']['mailto']
            servers_info = conf['servers']

            servers = []
            for server_info in servers_info:
                projects = []
                for project_info in server_info['projects']:
                    project = Server.Project(project_info['path'],
                                             project_info['port'])
                    projects.append(project)

                server = Server(
                    name=server_info['name'],
                    ip=server_info['ip'],
                    projects=projects,
                    passwd=server_info['passwd'],
                    default_passwd=default_password)

                servers.append(server)

            global SERVERS
            SERVERS = servers

            # initilize
            Server.Disk.set_threshold(disk_thresold)
            Email.setup(smtp_server, smtp_port, smtp_username, smtp_password,
                        mailto_list)
            disk_monitor = DiskMonitor(disk_monitor_interval)
            progress_monitor = ProgressMonitor(progress_monitor_interval)
            gateway_monitor = GatewayMonitor(gateway_monitor_interval)
            db_conn_monitor = DbConnMonitor(db_conn_monitor_interval)

            # clear cache
            disk_monitor.servers.clear()
            progress_monitor.servers.clear()
            gateway_monitor.servers.clear()
            # recache
            for server in servers:
                disk_monitor.servers.append(server)
                if server.any_project_has_port():
                    progress_monitor.servers.append(server)
                if 'gate' in server.name:
                    gateway_monitor.servers.append(server)
                if 'db' in server.name or 'datacenter' in server.name:
                    db_conn_monitor.servers.append(server)

            return conf, disk_monitor, progress_monitor, gateway_monitor, db_conn_monitor
        except Exception as e:
            logger.error('load config file error! details:\n{e}'.format(e=e))
            return None
        finally:
            file.close()
    else:
        logger.error('Cound not found config file!')
        return None
예제 #7
0
def summarize_amazing_race(path_to_data='sample_data/test1/', num_processes=4):
    """Calculates the amazing race statistics and prints them out for each person in the race.

    Goes through all the days annotated in index.json file. From each day it gatheres data for each participant and then
    combines data for whole race. The method can run in parallel, recommended number of processes is
    < number_cpu_cores * 2. While the data is being processed the progress is written to std in 5% increments by number
    of legs. At the end the data for each participant is printed to stdout.

    :param path_to_data: string, path to the folder containing data_xxx.json and index.json
    :param num_processes: int, number of processes to run in parallel
    """

    path = pathlib.Path(path_to_data)
    with open(
            path / "index.json", "r"
    ) as f:  # Because you use with, indent everything below to the same level
        index_jsonable = json.load(f)

        #  Read data from index.json file and stores it for later use.
        friend_list = index_jsonable['friends']
        total_distance = sum(dist for _, dist, _ in index_jsonable['files'])
        total_legs = sum(num_legs for *_, num_legs in index_jsonable['files'])
        file_names = [fn for fn, *_ in index_jsonable['files']]
        total_files = len(file_names)

        queue = mp.Queue(
        )  # Create shared queue in order to bring back results after they have been processed.
        progress_monitor = ProgressMonitor(
            total_legs, total_distance)  # Create shared ProgressMonitor

        processes = []
        step_size = math.ceil(total_files /
                              num_processes)  # How many files per process

        for i in range(0, len(file_names), step_size):
            #  Write all filename that are to be processed in this batch (by a single process)
            batch = [
                path / file_names[j]
                for j in range(i, min(i + step_size, total_files))
            ]
            processes.append(
                mp.Process(target=summarize_data_for_batch,
                           args=(batch, i // step_size, queue,
                                 progress_monitor)))

        for t in processes:
            t.start()
        for t in processes:
            t.join()

        #  Print final progress, it should be at 100%.
        progress_monitor.print_progress(True)

        #  Gather shared data from the queue.
        all_data = [None] * queue.qsize()
        while not queue.empty():
            processed_day = queue.get()
            all_data[processed_day[0]] = processed_day[1]

        friend_summaries = combine_data(
            all_data)  # Combine data from different batches into one.

        print("\n\n++++++++++++")
        print("In total %i friends participated in this race: " %
              (len(friend_list)) + ", ".join(friend_list))
        print("Together they traveled for total distance of %0.1f km\n" %
              total_distance)
        for friend in friend_summaries.keys():
            friend_summary = friend_summaries[friend]
            print("%s\n--------" % (friend, ))
            friend_summary.print()
            print("")
예제 #8
0
class StatusBar(gtk.VBox):
	
	STATUS_CODES = {
		-1  : "Communication Error", # STATUS_ERROR
		0x00: "Reset",               # STATUS_RESET
		0x01: "Device Not Ready",    # STATUS_BUSY
		0x40: "Stopped",             # STATUS_STOPPED
		0x41: "Breakpoint",          # STATUS_STOPPED_BREAKPOINT
		0x42: "Watchpoint",          # STATUS_STOPPED_WATCHPOINT
		0x43: "Memory Fault",        # STATUS_STOPPED_MEM_FAULT
		0x44: "Program Stopped",     # STATUS_STOPPED_PROG_REQ
		0x80: "Running",             # STATUS_RUNNING
		0x81: "Running (SWI)",       # STATUS_RUNNING_SWI
	}
	
	
	def __init__(self, system):
		"""
		A status bar for a system.
		"""
		gtk.VBox.__init__(self)
		
		self.system = system
		
		# Create a separator between us and the rest of the window
		sep = gtk.HSeparator()
		self.pack_start(sep, fill = True, expand = False)
		
		# A container for all the status entries
		self.hbox = gtk.HBox(spacing = 10)
		self.hbox.set_border_width(2)
		self.pack_start(self.hbox, fill = True, expand = False)
		
		# Current system status
		self.status_label = gtk.Label()
		self.hbox.pack_start(self.status_label, fill = True, expand = False)
		self._add_sep()
		
		# Step counter
		self.step_count_label = gtk.Label()
		self.hbox.pack_start(self.step_count_label, fill = True, expand = False)
		self._add_sep()
		
		# Show progress of background tasks
		self.progress_monitor = ProgressMonitor(self.system,
		                                        orientation = gtk.ORIENTATION_HORIZONTAL)
		self.hbox.pack_start(self.progress_monitor, fill = True, expand = True)
	
	
	def _add_sep(self):
		"""
		Add a horizontal separator to the end of the status bar
		"""
		sep = gtk.VSeparator()
		self.hbox.pack_start(sep, fill = True, expand = False)
	
	
	def add_adjustment(self, *args, **kwargs):
		# Forward to the progress_monitor
		self.progress_monitor.add_adjustment(*args, **kwargs)
	
	
	def remove_adjustment(self, *args, **kwargs):
		# Forward to the progress_monitor
		self.progress_monitor.remove_adjustment(*args, **kwargs)
	
	
	@RunInBackground()
	def refresh(self):
		status, steps_remaining, steps_since_reset = self.system.get_status()
		
		yield
		
		self.status_label.set_text(StatusBar.STATUS_CODES.get(status, "Device in Unknown State"))
		self.step_count_label.set_text("%s Step%s Since Reset%s"%(
			steps_since_reset,
			"s" if steps_since_reset != 1 else "",
			(" (%s Remaining)"%steps_remaining if steps_remaining > 0 else "")
		))
	
	
	def architecture_changed(self):
		"""
		Called when the architecture changes, deals with all the
		architecture-specific changes which need to be made to the GUI.
		"""
		# Nothing to do!
		self.refresh()
예제 #9
0
class Upscaler:
    """ An instance of this class is a upscaler that will
    upscale all images in the given directory.

    Raises:
        Exception -- all exceptions
        ArgumentError -- if argument is not valid
    """

    def __init__(self, input_path, output_path, driver_settings, ffmpeg_settings):
        # mandatory arguments
        self.input_path = input_path
        self.output_path = output_path
        self.driver_settings = driver_settings
        self.ffmpeg_settings = ffmpeg_settings

        # optional arguments
        self.driver = 'waifu2x_caffe'
        self.scale_width = None
        self.scale_height = None
        self.scale_ratio = None
        self.processes = 1
        self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
        self.image_format = 'png'
        self.preserve_frames = False

        # other internal members and signals
        self.stop_signal = False
        self.total_frames_upscaled = 0
        self.total_frames = 0

    def create_temp_directories(self):
        """create temporary directories
        """

        # if cache directory unspecified, use %TEMP%\video2x
        if self.video2x_cache_directory is None:
            self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'

        # if specified cache path exists and isn't a directory
        if self.video2x_cache_directory.exists() and not self.video2x_cache_directory.is_dir():
            Avalon.error(_('Specified or default cache directory is a file/link'))
            raise FileExistsError('Specified or default cache directory is a file/link')

        # if cache directory doesn't exist, try creating it
        if not self.video2x_cache_directory.exists():
            try:
                Avalon.debug_info(_('Creating cache directory {}').format(self.video2x_cache_directory))
                self.video2x_cache_directory.mkdir(parents=True, exist_ok=True)
            except Exception as exception:
                Avalon.error(_('Unable to create {}').format(self.video2x_cache_directory))
                raise exception

        # create temp directories for extracted frames and upscaled frames
        self.extracted_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(_('Extracted frames are being saved to: {}').format(self.extracted_frames))
        self.upscaled_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory))
        Avalon.debug_info(_('Upscaled frames are being saved to: {}').format(self.upscaled_frames))

    def cleanup_temp_directories(self):
        """delete temp directories when done
        """
        if not self.preserve_frames:
            for directory in [self.extracted_frames, self.upscaled_frames, self.video2x_cache_directory]:
                try:
                    # avalon framework cannot be used if python is shutting down
                    # therefore, plain print is used
                    print(_('Cleaning up cache directory: {}').format(directory))
                    shutil.rmtree(directory)
                except (OSError, FileNotFoundError):
                    print(_('Unable to delete: {}').format(directory))
                    traceback.print_exc()

    def _check_arguments(self):
        # if input is a file
        if self.input_path.is_file():
            if self.output_path.is_dir():
                Avalon.error(_('Input and output path type mismatch'))
                Avalon.error(_('Input is single file but output is directory'))
                raise ArgumentError('input output path type mismatch')
            if not re.search(r'.*\..*$', str(self.output_path)):
                Avalon.error(_('No suffix found in output file path'))
                Avalon.error(_('Suffix must be specified for FFmpeg'))
                raise ArgumentError('no output video suffix specified')

        # if input is a directory
        elif self.input_path.is_dir():
            if self.output_path.is_file():
                Avalon.error(_('Input and output path type mismatch'))
                Avalon.error(_('Input is directory but output is existing single file'))
                raise ArgumentError('input output path type mismatch')

        # if input is neither
        else:
            Avalon.error(_('Input path is neither a file nor a directory'))
            raise FileNotFoundError(f'{self.input_path} is neither file nor directory')

        # check Fmpeg settings
        ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path'])
        if not ((pathlib.Path(ffmpeg_path / 'ffmpeg.exe').is_file() and
                pathlib.Path(ffmpeg_path / 'ffprobe.exe').is_file()) or
                (pathlib.Path(ffmpeg_path / 'ffmpeg').is_file() and
                pathlib.Path(ffmpeg_path / 'ffprobe').is_file())):
            Avalon.error(_('FFmpeg or FFprobe cannot be found under the specified path'))
            Avalon.error(_('Please check the configuration file settings'))
            raise FileNotFoundError(self.ffmpeg_settings['ffmpeg_path'])

        # check if driver settings
        driver_settings = copy.deepcopy(self.driver_settings)
        driver_path = driver_settings.pop('path')

        # check if driver path exists
        if not (pathlib.Path(driver_path).is_file() or pathlib.Path(f'{driver_path}.exe').is_file()):
            Avalon.error(_('Specified driver executable directory doesn\'t exist'))
            Avalon.error(_('Please check the configuration file settings'))
            raise FileNotFoundError(driver_path)

        # parse driver arguments using driver's parser
        # the parser will throw AttributeError if argument doesn't satisfy constraints
        try:
            driver_arguments = []
            for key in driver_settings.keys():

                value = driver_settings[key]

                if value is None or value is False:
                    continue

                else:
                    if len(key) == 1:
                        driver_arguments.append(f'-{key}')
                    else:
                        driver_arguments.append(f'--{key}')
                    # true means key is an option
                    if value is not True:
                        driver_arguments.append(str(value))

            DriverWrapperMain = getattr(importlib.import_module(f'wrappers.{self.driver}'), 'WrapperMain')
            DriverWrapperMain.parse_arguments(driver_arguments)
        except AttributeError as e:
            Avalon.error(_('Failed to parse driver argument: {}').format(e.args[0]))
            raise e

    def _upscale_frames(self):
        """ Upscale video frames with waifu2x-caffe

        This function upscales all the frames extracted
        by ffmpeg using the waifu2x-caffe binary.

        Arguments:
            w2 {Waifu2x Object} -- initialized waifu2x object
        """

        # initialize waifu2x driver
        if self.driver not in AVAILABLE_DRIVERS:
            raise UnrecognizedDriverError(_('Unrecognized driver: {}').format(self.driver))

        # list all images in the extracted frames
        frames = [(self.extracted_frames / f) for f in self.extracted_frames.iterdir() if f.is_file]

        # if we have less images than processes,
        # create only the processes necessary
        if len(frames) < self.processes:
            self.processes = len(frames)

        # create a directory for each process and append directory
        # name into a list
        process_directories = []
        for process_id in range(self.processes):
            process_directory = self.extracted_frames / str(process_id)
            process_directories.append(process_directory)

            # delete old directories and create new directories
            if process_directory.is_dir():
                shutil.rmtree(process_directory)
            process_directory.mkdir(parents=True, exist_ok=True)

        # waifu2x-converter-cpp will perform multi-threading within its own process
        if self.driver == 'waifu2x_converter_cpp':
            process_directories = [self.extracted_frames]

        else:
            # evenly distribute images into each directory
            # until there is none left in the directory
            for image in frames:
                # move image
                image.rename(process_directories[0] / image.name)
                # rotate list
                process_directories = process_directories[-1:] + process_directories[:-1]

        # create threads and start them
        for process_directory in process_directories:

            DriverWrapperMain = getattr(importlib.import_module(f'wrappers.{self.driver}'), 'WrapperMain')
            driver = DriverWrapperMain(copy.deepcopy(self.driver_settings))

            # if the driver being used is waifu2x-caffe
            if self.driver == 'waifu2x_caffe':
                self.process_pool.append(driver.upscale(process_directory,
                                                         self.upscaled_frames,
                                                         self.scale_ratio,
                                                         self.scale_width,
                                                         self.scale_height,
                                                         self.image_format,
                                                         self.bit_depth))

            # if the driver being used is waifu2x-converter-cpp
            elif self.driver == 'waifu2x_converter_cpp':
                self.process_pool.append(driver.upscale(process_directory,
                                                         self.upscaled_frames,
                                                         self.scale_ratio,
                                                         self.processes,
                                                         self.image_format))

            # if the driver being used is waifu2x-ncnn-vulkan
            elif self.driver == 'waifu2x_ncnn_vulkan':
                self.process_pool.append(driver.upscale(process_directory,
                                                         self.upscaled_frames,
                                                         self.scale_ratio))

            # if the driver being used is srmd_ncnn_vulkan
            elif self.driver == 'srmd_ncnn_vulkan':
                self.process_pool.append(driver.upscale(process_directory,
                                                         self.upscaled_frames,
                                                         self.scale_ratio))

        # start progress bar in a different thread
        Avalon.debug_info(_('Starting progress monitor'))
        self.progress_monitor = ProgressMonitor(self, process_directories)
        self.progress_monitor.start()

        # create the clearer and start it
        Avalon.debug_info(_('Starting upscaled image cleaner'))
        self.image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(self.process_pool))
        self.image_cleaner.start()

        # wait for all process to exit
        try:
            self._wait()
        except (Exception, KeyboardInterrupt, SystemExit) as e:
            # cleanup
            Avalon.debug_info(_('Killing progress monitor'))
            self.progress_monitor.stop()

            Avalon.debug_info(_('Killing upscaled image cleaner'))
            self.image_cleaner.stop()
            raise e

        # if the driver is waifu2x-converter-cpp
        # images need to be renamed to be recognizable for FFmpeg
        if self.driver == 'waifu2x_converter_cpp':
            for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]:
                renamed = re.sub(f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.image_format}',
                                 f'.{self.image_format}',
                                 str(image.name))
                (self.upscaled_frames / image).rename(self.upscaled_frames / renamed)

        # upscaling done, kill helper threads
        Avalon.debug_info(_('Killing progress monitor'))
        self.progress_monitor.stop()

        Avalon.debug_info(_('Killing upscaled image cleaner'))
        self.image_cleaner.stop()

    def _terminate_subprocesses(self):
        Avalon.warning(_('Terminating all processes'))
        for process in self.process_pool:
            process.terminate()

    def _wait(self):
        """ wait for subprocesses in process pool to complete
        """
        Avalon.debug_info(_('Main process waiting for subprocesses to exit'))

        try:
            # while process pool not empty
            while self.process_pool:

                # if stop signal received, terminate all processes
                if self.stop_signal is True:
                    raise SystemExit

                for process in self.process_pool:
                    process_status = process.poll()

                    # if process finished
                    if process_status is None:
                        continue

                    # if return code is not 0
                    elif process_status != 0:
                        Avalon.error(_('Subprocess {} exited with code {}').format(process.pid, process_status))
                        raise subprocess.CalledProcessError(process_status, process.args)

                    else:
                        Avalon.debug_info(_('Subprocess {} exited with code {}').format(process.pid, process_status))
                        self.process_pool.remove(process)

                time.sleep(0.1)

        except (KeyboardInterrupt, SystemExit) as e:
            Avalon.warning(_('Stop signal received'))
            self._terminate_subprocesses()
            raise e

        except (Exception, subprocess.CalledProcessError) as e:
            Avalon.error(_('Subprocess execution ran into an error'))
            self._terminate_subprocesses()
            raise e

    def run(self):
        """ Main controller for Video2X

        This function controls the flow of video conversion
        and handles all necessary functions.
        """

        # external stop signal when called in a thread
        self.stop_signal = False

        # define process pool to contain processes
        self.process_pool = []

        # parse arguments for waifu2x
        # check argument sanity
        self._check_arguments()

        # define processing queue
        processing_queue = queue.Queue()

        # if input specified is single file
        if self.input_path.is_file():
            Avalon.info(_('Upscaling single video file: {}').format(self.input_path))
            processing_queue.put((self.input_path.absolute(), self.output_path.absolute()))

        # if input specified is a directory
        elif self.input_path.is_dir():

            # make output directory if it doesn't exist
            self.output_path.mkdir(parents=True, exist_ok=True)
            for input_video in [f for f in self.input_path.iterdir() if f.is_file()]:
                output_video = self.output_path / input_video.name
                processing_queue.put((input_video.absolute(), output_video.absolute()))

        while not processing_queue.empty():
            input_video, output_video = processing_queue.get()
            # drivers that have native support for video processing
            if self.driver == 'anime4kcpp':
                # append FFmpeg path to the end of PATH
                # Anime4KCPP will then use FFmpeg to migrate audio tracks
                os.environ['PATH'] += f';{self.ffmpeg_settings["ffmpeg_path"]}'
                Avalon.info(_('Starting to upscale extracted images'))

                # import and initialize Anime4KCPP wrapper
                DriverWrapperMain = getattr(importlib.import_module('wrappers.anime4kcpp'), 'WrapperMain')
                driver = DriverWrapperMain(copy.deepcopy(self.driver_settings))

                # run Anime4KCPP
                self.process_pool.append(driver.upscale(input_video, output_video, self.scale_ratio, self.processes))
                self._wait()
                Avalon.info(_('Upscaling completed'))

            else:
                try:
                    self.create_temp_directories()

                    # initialize objects for ffmpeg and waifu2x-caffe
                    fm = Ffmpeg(self.ffmpeg_settings, self.image_format)

                    Avalon.info(_('Reading video information'))
                    video_info = fm.get_video_info(input_video)
                    # analyze original video with ffprobe and retrieve framerate
                    # width, height = info['streams'][0]['width'], info['streams'][0]['height']

                    # find index of video stream
                    video_stream_index = None
                    for stream in video_info['streams']:
                        if stream['codec_type'] == 'video':
                            video_stream_index = stream['index']
                            break

                    # exit if no video stream found
                    if video_stream_index is None:
                        Avalon.error(_('Aborting: No video stream found'))
                        raise StreamNotFoundError('no video stream found')

                    # extract frames from video
                    self.process_pool.append((fm.extract_frames(input_video, self.extracted_frames)))
                    self._wait()

                    # get average frame rate of video stream
                    framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
                    fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']

                    # get a dict of all pixel formats and corresponding bit depth
                    pixel_formats = fm.get_pixel_formats()

                    # try getting pixel format's corresponding bti depth
                    try:
                        self.bit_depth = pixel_formats[fm.pixel_format]
                    except KeyError:
                        Avalon.error(_('Unsupported pixel format: {}').format(fm.pixel_format))
                        raise UnsupportedPixelError(f'unsupported pixel format {fm.pixel_format}')

                    Avalon.info(_('Framerate: {}').format(framerate))

                    # width/height will be coded width/height x upscale factor
                    if self.scale_ratio:
                        original_width = video_info['streams'][video_stream_index]['width']
                        original_height = video_info['streams'][video_stream_index]['height']
                        self.scale_width = int(self.scale_ratio * original_width)
                        self.scale_height = int(self.scale_ratio * original_height)

                    # upscale images one by one using waifu2x
                    Avalon.info(_('Starting to upscale extracted images'))
                    self._upscale_frames()
                    Avalon.info(_('Upscaling completed'))

                    # frames to Video
                    Avalon.info(_('Converting extracted frames into video'))

                    # use user defined output size
                    self.process_pool.append(fm.convert_video(framerate, f'{self.scale_width}x{self.scale_height}', self.upscaled_frames))
                    self._wait()
                    Avalon.info(_('Conversion completed'))

                    # migrate audio tracks and subtitles
                    Avalon.info(_('Migrating audio tracks and subtitles to upscaled video'))
                    self.process_pool.append(fm.migrate_audio_tracks_subtitles(input_video, output_video, self.upscaled_frames))
                    self._wait()

                    # destroy temp directories
                    self.cleanup_temp_directories()

                except (Exception, KeyboardInterrupt, SystemExit) as e:
                    with contextlib.suppress(ValueError):
                        self.cleanup_temp_directories()
                    raise e