V1.1 alpha: Fully working, added docs

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):
+    def extract_frames(self, videoin, outpath):
-        os.system("{} -i {} {}/extracted_%0d.png -y".format(self.ffmpeg_path, 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):
+        This method converts a set of images into a
-        os.system("{} -i {} -i {}/output-audio.aac -codec copy -shortest {} -y".format(self.ffmpeg_path, videoin, outpath, self.outfile))
+        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))

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"
+FOLDERIN = "frames"  # Folder containing extracted frames
-FOLDEROUT = "upscaled"
+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.
+    """Main controller for Video2X
-    This function takes care of the order of operation.
+
+    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()

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))