while True: fileName = Avalon.gets('请输入要保存的文件名,必须以.md为后缀:') if fileName.split('.')[-1] != 'md': Avalon.warning('文件名错误!') else: try: with open(fileName, 'w+', encoding='utf-8') as f: if not f.writable: raise Exception except: Avalon.warning('文件错误!') else: break Avalon.debug_info('程序已经启动...正在获取帖子信息') post = api(postID=postID, seeLZ=onlySeeLZ, debug=GENERAL_DEBUG_MODE) postInfo = post.getInfo() Avalon.time_info('开始任务:"%s"(作者:%s)' % (postInfo['Title'], postInfo['Author']), highlight=True) for i in range(1, postInfo['TotalPage'] + 1): Avalon.time_info('开始第%d页,共%d页' % (i, postInfo['TotalPage'])) try: pageHTMLContent = post.getContent(i) pageMarkdownContent = post.contentToMarkdown(pageHTMLContent, useImageBed=USE_IMAGE_BED) except KeyboardInterrupt: Avalon.critical('用户强制退出') quit(1) else: post.saveToFile(fileName, pageMarkdownContent)
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions): """This is the core function for WAIFU2X class Arguments: input_directory {string} -- source directory path output_directory {string} -- output directory path ratio {int} -- output video ratio """ try: # overwrite config file settings self.waifu2x_settings['input'] = input_directory self.waifu2x_settings['output'] = output_directory # print thread start message self.print_lock.acquire() Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started') self.print_lock.release() # waifu2x_ncnn_vulkan does not have long-opts, we'll have a dictionary that maps "our" config long-opt # names to their short opts waifu2x_ncnn_vulkan_opt_flag = { 'input': '-i', 'output': '-o', 'noise-level': '-n', 'scale-ratio': '-s', 'tile-size': '-t', 'model-path': '-m', 'gpu': '-g', 'load-proc-save_threads': '-j', 'verbose': '-v' } execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']] for key in self.waifu2x_settings.keys(): value = self.waifu2x_settings[key] if key == 'waifu2x_ncnn_vulkan_path': continue elif key == 'input': execute.append(waifu2x_ncnn_vulkan_opt_flag[key]) execute.append(input_directory) elif key == 'output': execute.append(waifu2x_ncnn_vulkan_opt_flag[key]) execute.append(output_directory) elif key == 'scale-ratio': execute.append(waifu2x_ncnn_vulkan_opt_flag[key]) # waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2 if scale_ratio == 1: execute.append('1') else: execute.append('2') # allow upper if cases to take precedence elif value is None or value is False: continue else: execute.append(waifu2x_ncnn_vulkan_opt_flag[key]) execute.append(str(value)) Avalon.debug_info(f'Executing: {execute}') subprocess.run(execute, check=True, stderr=subprocess.DEVNULL) # print thread exiting message self.print_lock.acquire() Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting') self.print_lock.release() return 0 except Exception as e: upscaler_exceptions.append(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
# resolve NTP server hostname into IP address # this makes sure that the hostname/IP is valid try: ntp_server_ip = str(netaddr.IPAddress(socket.gethostbyname(ntp_server))) except (socket.gaierror, netaddr.core.AddrFormatError) as error: Avalon.error(f"Unresolvable or invalid NTP hostname: {ntp_server}") raise error thread = threading.Thread(target=attack, args=(ntp_server_ip, args.target, args.count,)) thread.daemon = True thread.name = str(ntp_id) thread_pool.append(thread) Avalon.info('Launching attack') Avalon.info('Press Ctrl+C to stop the script') # start all threads after they have been added for thread in thread_pool: Avalon.debug_info(f"Starting thread {thread.name}") thread.start() # sleep and wait for SystemExit and KeyboardInterrupt try: for thread in thread_pool: thread.join() except (SystemExit, KeyboardInterrupt): Avalon.warning("Exiting signal received, script exiting") finally: Avalon.info("Script finished")
def generate_configs(output_path): """ Generate configuration file for every peer This function reads the PEERS list, generates a configuration file for every peer, and export into the CONFIG_OUTPUT directory. """ if len(pm.peers) == 0: Avalon.warning('No peers configured, exiting') exit(0) if len(pm.peers) == 1: Avalon.warning('Only one peer configured') Avalon.info('Generating configuration files') # Abort is destination is a file / link if os.path.isfile(output_path) or os.path.islink(output_path): Avalon.warning('Destination path is a file / link') Avalon.warning('Aborting configuration generation') return 1 # Ask if user wants to create the output directory if it doesn't exist if not os.path.isdir(output_path): if Avalon.ask( 'Output directory doesn\'t exist. Create output directory?', True): os.mkdir(output_path) else: Avalon.warning('Aborting configuration generation') return 1 # Iterate through all peers and generate configuration for each peer for peer in pm.peers: Avalon.debug_info('Generating configuration file for {}'.format( peer.address)) with open('{}/{}.conf'.format(output_path, peer.address.split('/')[0]), 'w') as config: # Write Interface configuration config.write('[Interface]\n') config.write('PrivateKey = {}\n'.format(peer.private_key)) if peer.address != '': config.write('Address = {}\n'.format(peer.address)) if peer.listen_port != '': config.write('ListenPort = {}\n'.format(peer.listen_port)) # Write peers' information for p in pm.peers: if p.address == peer.address: # Skip if peer is self continue config.write('\n[Peer]\n') print(p.private_key) config.write('PublicKey = {}\n'.format(wg.pubkey( p.private_key))) config.write('AllowedIPs = {}\n'.format(p.address)) if p.public_address != '': config.write('Endpoint = {}:{}\n'.format( p.public_address, p.listen_port)) if peer.keep_alive: config.write('PersistentKeepalive = 25\n') if p.preshared_key: config.write('PresharedKey = {}\n'.format(p.preshared_key))
def upscale(self, input_directory, output_directory, scale_ratio, scale_width, scale_height, image_format, upscaler_exceptions): """This is the core function for WAIFU2X class Arguments: input_directory {string} -- source directory path output_directory {string} -- output directory path width {int} -- output video width height {int} -- output video height """ try: # overwrite config file settings self.waifu2x_settings['input_path'] = input_directory self.waifu2x_settings['output_path'] = output_directory if scale_ratio: self.waifu2x_settings['scale_ratio'] = scale_ratio elif scale_width and scale_height: self.waifu2x_settings['scale_width'] = scale_width self.waifu2x_settings['scale_height'] = scale_height self.waifu2x_settings['output_extention'] = image_format # print thread start message self.print_lock.acquire() Avalon.debug_info( f'[upscaler] Thread {threading.current_thread().name} started') self.print_lock.release() # list to be executed execute = [] execute.append(self.waifu2x_settings['waifu2x_caffe_path']) for key in self.waifu2x_settings.keys(): value = self.waifu2x_settings[key] # is executable key or null or None means that leave this option out (keep default) if key == 'waifu2x_caffe_path' or value is None or value is False: continue else: if len(key) == 1: execute.append(f'-{key}') else: execute.append(f'--{key}') execute.append(str(value)) Avalon.debug_info(f'Executing: {execute}') completed_command = subprocess.run(execute, check=True) # print thread exiting message self.print_lock.acquire() Avalon.debug_info( f'[upscaler] Thread {threading.current_thread().name} exiting') self.print_lock.release() # return command execution return code return completed_command.returncode except Exception as e: upscaler_exceptions.append(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 """ # progress bar process exit signal self.progress_bar_exit_signal = False # initialize waifu2x driver if self.waifu2x_driver not in AVAILABLE_DRIVERS: raise UnrecognizedDriverError( f'Unrecognized driver: {self.waifu2x_driver}') # create a container for all upscaler processes upscaler_processes = [] # 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.waifu2x_driver in ['waifu2x_converter', 'anime4k']: 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: # if the driver being used is waifu2x-caffe if self.waifu2x_driver == 'waifu2x_caffe': driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings), self.method, self.model_dir, self.bit_depth) if self.scale_ratio: upscaler_processes.append( driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio, False, False, self.image_format)) else: upscaler_processes.append( driver.upscale(process_directory, self.upscaled_frames, False, self.scale_width, self.scale_height, self.image_format)) # if the driver being used is waifu2x-converter-cpp elif self.waifu2x_driver == 'waifu2x_converter': driver = Waifu2xConverter(self.driver_settings, self.model_dir) upscaler_processes.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.waifu2x_driver == 'waifu2x_ncnn_vulkan': driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings)) upscaler_processes.append( driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio)) # if the driver being used is anime4k elif self.waifu2x_driver == 'anime4k': driver = Anime4k(copy.deepcopy(self.driver_settings)) upscaler_processes += driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio, self.processes) # if the driver being used is srmd_ncnn_vulkan elif self.waifu2x_driver == 'srmd_ncnn_vulkan': driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings)) upscaler_processes.append( driver.upscale(process_directory, self.upscaled_frames, self.scale_ratio)) # start progress bar in a different thread progress_bar = threading.Thread(target=self._progress_bar, args=(process_directories, )) progress_bar.start() # create the clearer and start it Avalon.debug_info('Starting upscaled image cleaner') image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_processes)) image_cleaner.start() # wait for all process to exit try: Avalon.debug_info('Main process waiting for subprocesses to exit') for process in upscaler_processes: Avalon.debug_info( f'Subprocess {process.pid} exited with code {process.wait()}' ) except (KeyboardInterrupt, SystemExit): Avalon.warning('Exit signal received') Avalon.warning('Killing processes') for process in upscaler_processes: process.terminate() # cleanup and exit with exit code 1 Avalon.debug_info('Killing upscaled image cleaner') image_cleaner.stop() self.progress_bar_exit_signal = True sys.exit(1) # if the driver is waifu2x-converter-cpp # images need to be renamed to be recognizable for FFmpeg if self.waifu2x_driver == 'waifu2x_converter': 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 the clearer Avalon.debug_info('Killing upscaled image cleaner') image_cleaner.stop() # pass exit signal to progress bar thread self.progress_bar_exit_signal = True
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 upscale(self, input_directory, output_directory, scale_ratio, processes, push_strength=None, push_grad_strength=None): """ Anime4K wrapper Arguments: file_in {string} -- input file path file_out {string} -- output file path Keyword Arguments: scale {int} -- scale ratio (default: {None}) push_strength {int} -- residual push strength (default: {None}) push_grad_strength {int} -- residual gradient push strength (default: {None}) Returns: subprocess.Popen.returncode -- command line return value of execution """ # a list of all commands to be executed commands = queue.Queue() # get a list lof all image files in input_directory extracted_frame_files = [ f for f in input_directory.iterdir() if str(f).lower().endswith('.png') or str(f).lower().endswith('.jpg') ] # upscale each image in input_directory for image in extracted_frame_files: execute = [ self.driver_settings['java_path'], '-jar', self.driver_settings['path'], str(image.absolute()), str(output_directory / image.name), str(scale_ratio) ] # optional arguments kwargs = ['push_strength', 'push_grad_strength'] # if optional argument specified, append value to execution list for arg in kwargs: if locals()[arg] is not None: execute.extend([locals([arg])]) commands.put(execute) # initialize two lists to hold running and finished processes anime4k_running_processes = [] anime4k_finished_processes = [] # run all commands in queue while not commands.empty(): # if any commands have completed # remove the subprocess.Popen project and move it into finished processes for process in anime4k_running_processes: if process.poll() is not None: Avalon.debug_info( f'Subprocess {process.pid} exited with code {process.poll()}' ) anime4k_finished_processes.append(process) anime4k_running_processes.remove(process) # when number running processes is less than what's specified # create new processes and add to running process pool while len(anime4k_running_processes) < processes: next_in_queue = commands.get() new_process = subprocess.Popen(next_in_queue) anime4k_running_processes.append(new_process) self.print_lock.acquire() Avalon.debug_info( f'[upscaler] Subprocess {new_process.pid} executing: {shlex.join(next_in_queue)}' ) self.print_lock.release() # return command execution return code return anime4k_finished_processes
def _execute(self, execute): # turn all list elements into string to avoid errors execute = [str(e) for e in execute] Avalon.debug_info(f'Executing: {shlex.join(execute)}') return subprocess.Popen(execute)
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions): """This is the core function for WAIFU2X class Arguments: input_directory {string} -- source directory path output_directory {string} -- output directory path ratio {int} -- output video ratio """ try: # overwrite config file settings self.waifu2x_settings['input_path'] = input_directory self.waifu2x_settings['output_path'] = output_directory # print thread start message self.print_lock.acquire() Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started') self.print_lock.release() # waifu2x_ncnn_vulkan accepts arguments in a positional manner # See: https://github.com/nihui/waifu2x_ncnn_vulkan#usage # waifu2x_ncnn_vulkan.exe [input image] [output png] [noise=-1/0/1/2/3] [scale=1/2] [blocksize=400] # noise = noise level, large value means strong denoise effect, -1=no effect # scale = scale level, 1=no scale, 2=upscale 2x # blocksize = tile size, use smaller value to reduce GPU memory usage, default is 400 # waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2 if scale_ratio == 1: for raw_frame in os.listdir(input_directory): command = [ os.path.join(input_directory, raw_frame), os.path.join(output_directory, raw_frame), str(self.waifu2x_settings['noise-level']), '1', str(self.waifu2x_settings['block-size']) ] execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']] execute.extend(command) Avalon.debug_info(f'Executing: {execute}') subprocess.run(execute, check=True, stderr=subprocess.DEVNULL) else: for raw_frame in os.listdir(input_directory): command = [ os.path.join(input_directory, raw_frame), os.path.join(output_directory, raw_frame), str(self.waifu2x_settings['noise-level']), '2', str(self.waifu2x_settings['block-size']) ] execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']] execute.extend(command) Avalon.debug_info(f'Executing: {execute}') subprocess.run(execute, check=True, stderr=subprocess.DEVNULL) # print thread exiting message self.print_lock.acquire() Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting') self.print_lock.release() return 0 except Exception as e: upscaler_exceptions.append(e)
# -------------------- Execution if __name__ != "__main__": Avalon.error("This file cannot be imported as a library") raise ImportError("file cannot be imported") # parse command line arguments args = parse_arguments() # print KPM icon print_icon() try: # create KPM object kpm = Kpm() Avalon.debug_info("KPM initialized") # privileged section # check user privilege if os.getuid() != 0: Avalon.error("This program must be run as root") sys.exit(1) # if --install_kpm if args.install_kpm: # copy the current file to defined binary path shutil.copy(pathlib.Path(__file__), KPM_PATH) # change owner and permission of the file os.chown(KPM_PATH, 0, 0) KPM_PATH.chmod(0o755)
# print video2x logo print_logo() # parse command line arguments video2x_args, driver_args = parse_arguments() # display version and lawful informaition if video2x_args.version: print(LEGAL_INFO) sys.exit(0) # redirect output to both terminal and log file if video2x_args.disable_logging is False: LOGFILE = video2x_args.log Avalon.debug_info(_('Redirecting console logs to {}').format(LOGFILE)) sys.stdout = BiLogger(sys.stdout, LOGFILE) sys.stderr = BiLogger(sys.stderr, LOGFILE) # read configurations from configuration file config = read_config(video2x_args.config) # load waifu2x configuration driver_settings = config[video2x_args.driver] driver_settings['path'] = os.path.expandvars(driver_settings['path']) # read FFmpeg configuration ffmpeg_settings = config['ffmpeg'] ffmpeg_settings['ffmpeg_path'] = os.path.expandvars( ffmpeg_settings['ffmpeg_path'])
if args.force_upgrade: Avalon.info('Force upgrading KPM from GitHub') upgrade_kpm() exit(0) check_version() # if -x, --xinstall specified if args.xinstall: packages = args.xinstall.split(',') kobj.xinstall(packages) # if no arguments are given else: kobj.upgrade_all() Avalon.debug_info('Checking for unused packages') # check if there are any unused packages if kobj.autoremove_available(): if Avalon.ask('Remove useless packages?', True): kobj.autoremove() else: Avalon.info('No unused packages found') # apt autoclean Avalon.info('Erasing old downloaded archive files') kobj.autoclean() except KeyboardInterrupt: Avalon.warning('Aborting')
def _upscale_frames(self, w2): """ 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 """ # progress bar thread exit signal self.progress_bar_exit_signal = False # create a container for exceptions in threads # if this thread is not empty, then an exception has occured self.upscaler_exceptions = [] # it's easier to do multi-threading with waifu2x_converter # the number of threads can be passed directly to waifu2x_converter if self.waifu2x_driver == 'waifu2x_converter': progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames], )) progress_bar.start() w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.upscaler_exceptions) for image in [ f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f)) ]: renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.png', '.png', image) shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed)) self.progress_bar_exit_signal = True progress_bar.join() return # create a container for all upscaler threads upscaler_threads = [] # list all images in the extracted frames frames = [ os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f)) ] # if we have less images than threads, # create only the threads necessary if len(frames) < self.threads: self.threads = len(frames) # create a folder for each thread and append folder # name into a list thread_pool = [] thread_folders = [] for thread_id in range(self.threads): thread_folder = '{}\\{}'.format(self.extracted_frames, str(thread_id)) thread_folders.append(thread_folder) # delete old folders and create new folders if os.path.isdir(thread_folder): shutil.rmtree(thread_folder) os.mkdir(thread_folder) # append folder path into list thread_pool.append((thread_folder, thread_id)) # evenly distribute images into each folder # until there is none left in the folder for image in frames: # move image shutil.move(image, thread_pool[0][0]) # rotate list thread_pool = thread_pool[-1:] + thread_pool[:-1] # create threads and start them for thread_info in thread_pool: # create thread if self.scale_ratio: thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.upscaler_exceptions)) else: thread = threading.Thread( target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.upscaler_exceptions)) thread.name = thread_info[1] # add threads into the pool upscaler_threads.append(thread) # start progress bar in a different thread progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders, )) progress_bar.start() # create the clearer and start it Avalon.debug_info('Starting upscaled image cleaner') image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads)) image_cleaner.start() # start all threads for thread in upscaler_threads: thread.start() # wait for threads to finish for thread in upscaler_threads: thread.join() # upscaling done, kill the clearer Avalon.debug_info('Killing upscaled image cleaner') image_cleaner.stop() self.progress_bar_exit_signal = True if len(self.upscaler_exceptions) != 0: raise (self.upscaler_exceptions[0])
def upgrade_all(self): """ upgrade all packages This method checks if there are packages available for updating and update the packages if updating them won't remove any packages from the system. Often times when a bad source is added to the system, APT tends to remove a number of packages from the system when upgrading which is very risky. """ Avalon.info('Starting automatic upgrade') Avalon.info('Updating APT cache') with open('/etc/apt/sources.list', 'r') as aptlist: for line in aptlist: if 'ubuntu.com' in line and distro.linux_distribution( )[0] != 'Ubuntu' and line.replace(' ', '')[0] != '#': Avalon.warning('Ubuntu source detected in source.list!') Avalon.warning( 'Continue upgrading might cause severe consequences!') if Avalon.ask('Are you sure that you want to continue?', False): break else: Avalon.warning('Aborting system upgrade..') exit(0) self.update() Avalon.info('APT cache updated') if len(self.import_list) != 0: if Avalon.ask('Detected unimported keys, import?', True): if shutil.which('dirmngr') is None: Avalon.warning('dirmngr Not installed') Avalon.warning('It is required for importing keys') # ask if user wants to install dirmngr if Avalon.ask('Install Now?'): self.install('dirnmgr') # check dirmngr after package installation if isinstance(shutil.which('dirmngr'), str): Avalon.info('Installation successful') self.import_keys(self.import_list) Avalon.info('Keys imported') Avalon.info( 'Updating APT cache after key importing') self.update() else: Avalon.error('Installation Failed') Avalon.error('Please check your settings') Avalon.warning( 'dirmngr not available. Continuing without importing keys' ) else: Avalon.warning('dirmngr not available') Avalon.warning('Continuing without importing keys') else: self.import_keys(self.import_list) Avalon.info('Keys imported') Avalon.info('Updating APT cache after key importing') self.update() # Second update after keys are imported self.update() # if there are no upgrades available Avalon.debug_info('Checking package updates') if self.no_upgrades(): Avalon.info('No upgrades available') # if upgrades are available else: Avalon.debug_info('Checking if full upgrade is safe') # if upgrade is safe, use -y flag on apt-get full-upgrade # otherwise, let user confirm the upgrade if self.full_upgrade_safe(): Avalon.info('Full upgrade is safe') Avalon.info('Starting APT full upgrade') self.full_upgrade() else: Avalon.warning('Full upgrade is NOT safe') Avalon.warning('Requiring human confirmation') self.manual_full_upgrade()
def wrapper(*args, **kwargs): return function(*args, **kwargs) Avalon.debug_info('{} row(s) affected'.format(args[0].cursor.rowcount))
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
def send_request(self): """send POST request to CloudFlare's server""" global successes global fails global proxies global thread_pool succeeded = False error_message = "" try: # generate random installation ID install_id = "".join( random.choices(string.ascii_uppercase + string.digits, k=22)) # construct POST request payload post_data = { "key": f"{''.join(random.choices(string.ascii_uppercase + string.digits, k=43))}=", "install_id": install_id, "fcm_token": f"{install_id}:APA91b{''.join(random.choices(string.ascii_uppercase + string.digits, k=134))}", "referrer": self.warpid, "warp_enabled": False, "tos": datetime.datetime.utcnow().isoformat()[:-3] + "+00:00", "type": "Android", "locale": "en_US", } # send request if proxies is None: response = requests.post( f"https://api.cloudflareclient.com/v0a{random.randint(100, 999)}/reg", json=post_data, headers=HEADERS, timeout=self.timeout, ) else: proxy = proxies.popleft() response = requests.post( f"https://api.cloudflareclient.com/v0a{random.randint(100, 999)}/reg", json=post_data, headers=HEADERS, timeout=self.timeout, proxies={ "http": f"socks4://{proxy}", "https": f"socks4://{proxy}", }, ) if response.status_code == requests.codes.ok: successes += 1 succeeded = True else: fails += 1 error_message = f" with code {response.status_code}" if proxies: proxies.append(proxy) except IndexError: self.running = False return # print error and carryon on upon exceptions except Exception as e: # _print(f"Thread {self.name} encountered an exception") # _print(traceback.format_exc(), file=sys.stderr) fails += 1 error_message = f" with Python exception {e}" if proxies and not self.autoremove: proxies.append(proxy) finally: information = [ f"Successes: {Avalon.FG.G}{successes}{Avalon.FG.DGR}", f"Fails: {Avalon.FG.R}{fails}{Avalon.FG.DGR}", # f"Live Threads: {threading.active_count()}", f"Live Threads: {Avalon.FG.W}{len([t for t in thread_pool if t.is_alive()])}{Avalon.FG.DGR}", ] # this information is not accurate # information.append(f"Proxies in Pool: {len(proxies)}") if proxies else None if succeeded: Avalon.debug_info( f"[{' | '.join(information)}] Thread {self.name} succeeded", ) else: Avalon.debug_info( f"[{' | '.join(information)}] Thread {self.name} failed{error_message}", )
def upscale(self, input_directory, output_directory, scale_ratio, jobs, image_format, upscaler_exceptions): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. Arguments: input_directory {string} -- source directory path output_directory {string} -- output directory path scale_ratio {int} -- frames' scale ratio threads {int} -- number of threads """ try: # overwrite config file settings self.waifu2x_settings['input'] = input_directory self.waifu2x_settings['output'] = output_directory # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 """ self.waifu2x_settings['i'] = input_directory self.waifu2x_settings['o'] = output_directory self.waifu2x_settings['input'] = None self.waifu2x_settings['output'] = None """ self.waifu2x_settings['scale-ratio'] = scale_ratio self.waifu2x_settings['jobs'] = jobs self.waifu2x_settings['output-format'] = image_format # models_rgb must be specified manually for waifu2x-converter-cpp # if it's not specified in the arguments, create automatically if self.waifu2x_settings['model-dir'] is None: self.waifu2x_settings['model-dir'] = pathlib.Path( self.waifu2x_settings['waifu2x_converter_path'] ) / 'models_rgb' # print thread start message self.print_lock.acquire() Avalon.debug_info( f'[upscaler] Thread {threading.current_thread().name} started') self.print_lock.release() # list to be executed execute = [] for key in self.waifu2x_settings.keys(): value = self.waifu2x_settings[key] # the key doesn't need to be passed in this case if key == 'waifu2x_converter_path': execute.append( pathlib.Path(str(value)) / 'waifu2x-converter-cpp.exe') # null or None means that leave this option out (keep default) elif value is None or value is False: continue else: if len(key) == 1: execute.append(f'-{key}') else: execute.append(f'--{key}') # true means key is an option if value is True: continue execute.append(str(value)) Avalon.debug_info(f'Executing: {execute}') return subprocess.run(execute, check=True).returncode except Exception as e: upscaler_exceptions.append(e)
def upscale(self, input_directory, output_directory, scale_ratio, jobs, image_format): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. Arguments: input_directory {string} -- source directory path output_directory {string} -- output directory path scale_ratio {int} -- frames' scale ratio threads {int} -- number of threads """ # overwrite config file settings self.driver_settings['input'] = input_directory self.driver_settings['output'] = output_directory self.driver_settings['scale-ratio'] = scale_ratio self.driver_settings['jobs'] = jobs self.driver_settings['output-format'] = image_format # models_rgb must be specified manually for waifu2x-converter-cpp # if it's not specified in the arguments, create automatically if self.driver_settings['model-dir'] is None: self.driver_settings['model-dir'] = pathlib.Path( self.driver_settings['path']) / 'models_rgb' # list to be executed # initialize the list with waifu2x binary path as the first element execute = [ str( pathlib.Path(self.driver_settings['path']) / 'waifu2x-converter-cpp') ] for key in self.driver_settings.keys(): value = self.driver_settings[key] # the key doesn't need to be passed in this case if key == 'path': continue # null or None means that leave this option out (keep default) elif value is None or value is False: continue else: if len(key) == 1: execute.append(f'-{key}') else: execute.append(f'--{key}') # true means key is an option if value is True: continue execute.append(str(value)) # return the Popen object of the new process created self.print_lock.acquire() Avalon.debug_info( f'[upscaler] Subprocess {os.getpid()} executing: {shlex.join(execute)}' ) self.print_lock.release() return subprocess.Popen(execute)
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 upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions, push_strength=None, push_grad_strength=None): """ Anime4K wrapper Arguments: file_in {string} -- input file path file_out {string} -- output file path Keyword Arguments: scale {int} -- scale ratio (default: {None}) push_strength {int} -- residual push strength (default: {None}) push_grad_strength {int} -- residual gradient push strength (default: {None}) Returns: subprocess.Popen.returncode -- command line return value of execution """ try: # return value is the sum of all execution return codes return_value = 0 # get a list lof all image files in input_directory extracted_frame_files = [f for f in input_directory.iterdir() if str(f).lower().endswith('.png') or str(f).lower().endswith('.jpg')] # Only print debug_info on first iteration. logged = False # upscale each image in input_directory for image in extracted_frame_files: execute = [ self.settings['java_path'] / self.settings['java_binary'], '-jar', self.settings['path'] / self.settings['binary'], image.absolute(), output_directory / image.name, scale_ratio ] # optional arguments kwargs = [ 'push_strength', 'push_grad_strength' ] # if optional argument specified, append value to execution list for arg in kwargs: if locals()[arg] is not None: execute.extend([locals([arg])]) # turn all list elements into string to avoid errors execute = [str(e) for e in execute] if not logged: self.print_lock.acquire() Avalon.debug_info(f'Executing: {execute}') self.print_lock.release() logged = True return_value += subprocess.run(execute, check=True).returncode # print thread exiting message self.print_lock.acquire() Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting') self.print_lock.release() # return command execution return code return return_value except Exception as e: upscaler_exceptions.append(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 """ # progress bar thread exit signal self.progress_bar_exit_signal = False # create a container for exceptions in threads # if this thread is not empty, then an exception has occured self.upscaler_exceptions = [] # initialize waifu2x driver drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'] if self.waifu2x_driver not in drivers: raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}') # it's easier to do multi-threading with waifu2x_converter # the number of threads can be passed directly to waifu2x_converter if self.waifu2x_driver == 'waifu2x_converter': w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir) progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],)) progress_bar.start() w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions) for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]: renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image) shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed)) self.progress_bar_exit_signal = True progress_bar.join() return else: # create a container for all upscaler threads upscaler_threads = [] # list all images in the extracted frames frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))] # if we have less images than threads, # create only the threads necessary if len(frames) < self.threads: self.threads = len(frames) # create a directory for each thread and append directory # name into a list thread_pool = [] thread_directories = [] for thread_id in range(self.threads): thread_directory = os.path.join(self.extracted_frames, str(thread_id)) thread_directories.append(thread_directory) # delete old directories and create new directories if os.path.isdir(thread_directory): shutil.rmtree(thread_directory) os.mkdir(thread_directory) # append directory path into list thread_pool.append((thread_directory, thread_id)) # evenly distribute images into each directory # until there is none left in the directory for image in frames: # move image shutil.move(image, thread_pool[0][0]) # rotate list thread_pool = thread_pool[-1:] + thread_pool[:-1] # create threads and start them for thread_info in thread_pool: # create a separate w2 instance for each thread if self.waifu2x_driver == 'waifu2x_caffe': w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir, self.bit_depth) if self.scale_ratio: thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.image_format, self.upscaler_exceptions)) else: thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.image_format, self.upscaler_exceptions)) # if the driver being used is waifu2x_ncnn_vulkan elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan': w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, self.upscaler_exceptions)) # create thread thread.name = thread_info[1] # add threads into the pool upscaler_threads.append(thread) # start progress bar in a different thread progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,)) progress_bar.start() # create the clearer and start it Avalon.debug_info('Starting upscaled image cleaner') image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads)) image_cleaner.start() # start all threads for thread in upscaler_threads: thread.start() # wait for threads to finish for thread in upscaler_threads: thread.join() # upscaling done, kill the clearer Avalon.debug_info('Killing upscaled image cleaner') image_cleaner.stop() self.progress_bar_exit_signal = True if len(self.upscaler_exceptions) != 0: raise(self.upscaler_exceptions[0])
# load cache directory if config['video2x']['video2x_cache_directory'] is not None: video2x_cache_directory = pathlib.Path( config['video2x']['video2x_cache_directory']) else: video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x' if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir(): Avalon.error(_('Specified cache directory is a file/link')) raise FileExistsError('Specified cache directory is a file/link') # if cache directory doesn't exist # ask the user if it should be created elif not video2x_cache_directory.exists(): try: Avalon.debug_info( _('Creating cache directory {}').format(video2x_cache_directory)) video2x_cache_directory.mkdir(parents=True, exist_ok=True) # there can be a number of exceptions here # PermissionError, FileExistsError, etc. # therefore, we put a catch-them-all here except Exception as exception: Avalon.error(_('Unable to create {}').format(video2x_cache_directory)) raise exception # start execution try: # start timer begin_time = time.time() # if input specified is a single file if video2x_args.input.is_file():