V1.1 alpha: Fully working, added docs

This commit is contained in:
K4YT3X 2018-02-24 22:52:04 -05:00
parent e198972c35
commit 994a542d52
3 changed files with 128 additions and 24 deletions

View File

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

View File

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

View File

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