basic functions online

ffmpeg.py

@@ -17,17 +17,18 @@ import os
 
 class FFMPEG:
 
-    def __init__(self):
+    def __init__(self, ffmpeg_path, outfile):
-        pass
+        self.ffmpeg_path = ffmpeg_path
+        self.outfile = outfile
 
     def strip_frames(self, videoin, outpath):
-        os.system("ffmpeg -i {} -r 1/1 {}/extracted_%0d.png".format(videoin, outpath))
+        os.system("{} -i {} {}/extracted_%0d.png -y".format(self.ffmpeg_path, videoin, outpath))
 
     def extract_audio(self, videoin, outpath):
-        os.system("ffmpeg -i {} -vn -acodec copy {}/output-audio.aac".format(videoin, outpath))
+        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("ffmpeg -r {} -f image2 -s {} -i {}/extracted_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p output.mp4".format(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))
 
-    def pressin_audio(self, videoin):
+    def pressin_audio(self, videoin, outpath):
-        os.system("ffmpeg -i {} -i audio.mp3 -codec copy -shortest output.mp4")
+        os.system("{} -i {} -i {}/output-audio.aac -codec copy -shortest {} -y".format(self.ffmpeg_path, videoin, outpath, self.outfile))

video2x.py

@@ -12,15 +12,76 @@ Version 1.0
 """
 
 from ffmpeg import FFMPEG
+from fractions import Fraction
 from waifu2x import WAIFU2X
+import argparse
+import json
 import os
 
-fm = FFMPEG()
+FFMPEG_PATH = "C:/Program Files (x86)/ffmpeg/bin/"
-w2 = WAIFU2X()
+WAIFU2X_PATH = "\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\""
-if not os.path.isdir("frames"):
-    os.mkdir("frames")
-fm.strip_frames("testf.mp4", "frames")
 
-if not os.path.isdir("upscaled"):
+FOLDERIN = "frames"
-    os.mkdir("upscaled")
+FOLDEROUT = "upscaled"
-w2.upscale("frames", "upscaled")
+
+
+def processArguments():
+    """This function parses all arguments
+    This allows users to customize options
+    for the output video.
+    """
+    global args
+    parser = argparse.ArgumentParser()
+    control_group = parser.add_argument_group('Controls')
+    control_group.add_argument("-f", "--factor", help="Factor to enlarge video by", action="store", default=2)
+    control_group.add_argument("-v", "--video", help="Specify video file", action="store", default=False)
+    control_group.add_argument("-o", "--output", help="Specify output file", action="store", default=False)
+    args = parser.parse_args()
+
+
+def get_vid_info():
+    """Gets original video information
+    This function uses ffprobe to determine the
+    properties of the original video.
+
+    It returns a dict
+    """
+    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')
+    json_str = json_file.read()
+    print(json.loads(json_str))
+    return json.loads(json_str)
+
+
+def main():
+    """Main flow control function for video2x.
+    This function takes care of the order of operation.
+    """
+    fm = FFMPEG("\"" + FFMPEG_PATH + "ffmpeg.exe\"", args.output)
+    w2 = WAIFU2X(WAIFU2X_PATH)
+
+    # Extract Frames
+    if not os.path.isdir(FOLDERIN):
+        os.mkdir(FOLDERIN)
+    fm.strip_frames(args.video, FOLDERIN)
+
+    info = get_vid_info()
+    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))
+
+    # Upscale Frames
+    if not os.path.isdir(FOLDEROUT):
+        os.mkdir(FOLDEROUT)
+    w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height)
+
+    # Frames to Video
+    fm.to_vid(framerate, final_resolution, FOLDEROUT)
+
+    # Extract and press audio in
+    fm.extract_audio(args.video, FOLDEROUT)
+    fm.pressin_audio("output.mp4", FOLDEROUT)
+
+
+processArguments()
+main()

waifu2x.py

@@ -17,8 +17,8 @@ import os
 
 class WAIFU2X:
 
-    def __init__(self):
+    def __init__(self, waifu2x_path):
-        pass
+        self.waifu2x_path = waifu2x_path
 
-    def upscale(self, factor, folderin, folderout):
+    def upscale(self, folderin, folderout, width, height):
-        os.system("waifu2x-caffe-cui.exe -i {} -o {}".format(folderin, folderout))
+        os.system("{} -p cpu -I png -i {} -e png -o {} -w {} -h {}".format(self.waifu2x_path, folderin, folderout, width, height))