From 994a542d5279bb201fe2bd8c3ace2780b73c632a Mon Sep 17 00:00:00 2001 From: K4YT3X Date: Sat, 24 Feb 2018 22:52:04 -0500 Subject: [PATCH] V1.1 alpha: Fully working, added docs --- ffmpeg.py | 61 ++++++++++++++++++++++++++++++++++++++++++++------ video2x.py | 65 ++++++++++++++++++++++++++++++++++++++++++------------ waifu2x.py | 26 +++++++++++++++++++--- 3 files changed, 128 insertions(+), 24 deletions(-) diff --git a/ffmpeg.py b/ffmpeg.py index a3c4d55..450a7e0 100644 --- a/ffmpeg.py +++ b/ffmpeg.py @@ -9,26 +9,73 @@ Last Modified: Feb 24, 2018 Description: This class handles all FFMPEG related operations. -Version 1.0 +Version 1.1 """ import os class FFMPEG: + """This class communicates with ffmpeg + + This class deals with ffmpeg. It handles extracitng + frames, stripping audio, converting images into videos + and inserting audio tracks to videos. + """ def __init__(self, ffmpeg_path, outfile): self.ffmpeg_path = ffmpeg_path self.outfile = outfile - def strip_frames(self, videoin, outpath): - os.system("{} -i {} {}/extracted_%0d.png -y".format(self.ffmpeg_path, videoin, outpath)) + def extract_frames(self, videoin, outpath): + """Extract every frame from original videos + + This method extracts every frame from videoin + using ffmpeg + + Arguments: + videoin {string} -- input video path + outpath {string} -- video output folder + """ + os.system( + "{} -i {} {}/extracted_%0d.png -y".format(self.ffmpeg_path, videoin, outpath)) def extract_audio(self, videoin, outpath): - os.system("{} -i {} -vn -acodec copy {}/output-audio.aac -y".format(self.ffmpeg_path, videoin, outpath)) + """Strips audio tracks from videos + + This method strips audio tracks from videos + into the output folder in aac format. + + Arguments: + videoin {string} -- input video path + outpath {string} -- video output folder + """ + os.system( + "{} -i {} -vn -acodec copy {}/output-audio.aac -y".format(self.ffmpeg_path, videoin, outpath)) def to_vid(self, framerate, resolution, folder): - os.system("{} -r {} -f image2 -s {} -i {}/extracted_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p output.mp4 -y".format(self.ffmpeg_path, framerate, resolution, folder)) + """Converts images into videos - def pressin_audio(self, videoin, outpath): - os.system("{} -i {} -i {}/output-audio.aac -codec copy -shortest {} -y".format(self.ffmpeg_path, videoin, outpath, self.outfile)) + This method converts a set of images into a + video. + + Arguments: + framerate {float} -- target video framerate + resolution {string} -- target video resolution + folder {string} -- source images folder + """ + os.system("{} -r {} -f image2 -s {} -i {}/extracted_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p output.mp4 -y".format( + self.ffmpeg_path, framerate, resolution, folder)) + + def insert_audio_track(self, videoin, outpath): + """Insert audio into video + + Inserts the AAC audio track stripped from + the original video into final video. + + Arguments: + videoin {string} -- input video path + outpath {string} -- video output folder + """ + os.system("{} -i {} -i {}/output-audio.aac -codec copy -shortest {} -y".format( + self.ffmpeg_path, videoin, outpath, self.outfile)) diff --git a/video2x.py b/video2x.py index 282fb34..505f884 100644 --- a/video2x.py +++ b/video2x.py @@ -1,14 +1,31 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ + +__ __ _ _ ___ __ __ +\ \ / / (_) | | |__ \ \ \ / / + \ \ / / _ __| | ___ ___ ) | \ V / + \ \/ / | | / _` | / _ \ / _ \ / / > < + \ / | | | (_| | | __/ | (_) | / /_ / . \ + \/ |_| \__,_| \___| \___/ |____| /_/ \_\ + + Name: Video2x Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: Feb 24, 2018 +Last Modified: Feb 25, 2018 -Description: This is the main controller for Video2x +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt -Version 1.0 +(C) 2016 - 2017 K4YT3X + +Description: Video2X is an automation software based on +waifu2x image enlarging engine. It extracts frames from a +video, enlarge it by a number of times without losing any +details or quality, keeping lines smooth and edges sharp. + +Version 1.1 alpha """ from ffmpeg import FFMPEG @@ -17,16 +34,21 @@ from waifu2x import WAIFU2X import argparse import json import os +import traceback +# FFMPEG bin folder. Mind that "/" at the end FFMPEG_PATH = "C:/Program Files (x86)/ffmpeg/bin/" +# waifu2x executable path. Mind all the forward slashes WAIFU2X_PATH = "\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\"" -FOLDERIN = "frames" -FOLDEROUT = "upscaled" +FOLDERIN = "frames" # Folder containing extracted frames +FOLDEROUT = "upscaled" # Folder contaning enlarges frames def processArguments(): - """This function parses all arguments + """Processes CLI arguments + + This function parses all arguments This allows users to customize options for the output video. """ @@ -44,10 +66,14 @@ def processArguments(): def get_vid_info(): """Gets original video information - This function uses ffprobe to determine the - properties of the original video. - It returns a dict + Retrieves original video information using + ffprobe, then export it into json file. + Finally it reads, parses the json file and + returns a dictionary + + Returns: + dictionary -- original video information """ os.system("{} -v quiet -print_format json -show_format -show_streams {} > info.json".format("\"" + FFMPEG_PATH + "ffprobe.exe\"", args.video)) json_file = open('info.json', 'r') @@ -57,8 +83,10 @@ def get_vid_info(): def main(): - """Main flow control function for video2x. - This function takes care of the order of operation. + """Main controller for Video2X + + This function controls the flow of video conversion + and handles all necessary functions. """ if args.cpu: method = "cpu" @@ -73,9 +101,10 @@ def main(): # Extract Frames if not os.path.isdir(FOLDERIN): os.mkdir(FOLDERIN) - fm.strip_frames(args.video, FOLDERIN) + fm.extract_frames(args.video, FOLDERIN) info = get_vid_info() + # Framerate is read as fraction from the json dictionary width, height, framerate = info["streams"][0]["width"], info["streams"][0]["height"], float(Fraction(info["streams"][0]["avg_frame_rate"])) print("Framerate: ", framerate) final_resolution = str(width * int(args.factor)) + "x" + str(height * int(args.factor)) @@ -90,11 +119,13 @@ def main(): # Extract and press audio in fm.extract_audio(args.video, FOLDEROUT) - fm.pressin_audio("output.mp4", FOLDEROUT) + fm.insert_audio_track("output.mp4", FOLDEROUT) processArguments() +# Check if arguments are valid / all necessary argument +# values are specified if not args.video: print("Error: You need to specify the video to process") exit(1) @@ -105,4 +136,10 @@ elif not args.cpu and not args.gpu and not args.cudnn: print("Error: You need to specify the enlarging processing unit") exit(1) -main() +if __name__ == '__main__': + try: + main() + except Exception as e: + # This code block is reserved for future + # fail-safe handlers + traceback.print_exc() diff --git a/waifu2x.py b/waifu2x.py index 730552f..2f923e1 100644 --- a/waifu2x.py +++ b/waifu2x.py @@ -4,22 +4,42 @@ Name: FFMPEG Class Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: Feb 24, 2018 +Last Modified: Feb 25, 2018 Description: This class controls waifu2x engine -Version 1.0 +Version 1.1 """ import os class WAIFU2X: + """This class communicates with waifu2x cui engine + + An object will be created for this class, containing information + about the binary address and the processing method. When being called + by the main program, other detailed information will be passed to + the upscale function. + + TODO: Make enhancement model customizable + """ def __init__(self, waifu2x_path, method): self.waifu2x_path = waifu2x_path self.method = method def upscale(self, folderin, folderout, width, height): - os.system("{} -p {} -I png -i {} -e png -o {} -w {} -h {} -n 3 -m noise_scale -y photo".format(self.waifu2x_path, self.method, folderin, folderout, width, height)) + """This is the core function for WAIFU2X class + + [description] + + Arguments: + folderin {string} -- source folder path + folderout {string} -- output folder path + width {int} -- output video width + height {int} -- output video height + """ + os.system("{} -p {} -I png -i {} -e png -o {} -w {} -h {} -n 3 -m noise_scale -y photo".format( + self.waifu2x_path, self.method, folderin, folderout, width, height))