video2x/video2x.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

__      __  _       _                  ___   __   __
\ \    / / (_)     | |                |__ \  \ \ / /
 \ \  / /   _    __| |   ___    ___      ) |  \ V /
  \ \/ /   | |  / _` |  / _ \  / _ \    / /    > <
   \  /    | | | (_| | |  __/ | (_) |  / /_   / . \
    \/     |_|  \__,_|  \___|  \___/  |____| /_/ \_\


Name: Video2x Controller
Author: K4YT3X
Date Created: Feb 24, 2018
Last Modified: May 19, 2018

Licensed under the GNU General Public License Version 3 (GNU GPL v3),
    available at: https://www.gnu.org/licenses/gpl-3.0.txt

(C) 2018 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.
"""
from avalon_framework import Avalon
from ffmpeg import FFMPEG
from fractions import Fraction
from waifu2x import WAIFU2X
import argparse
import inspect
import json
import os
import shutil
import subprocess
import traceback

VERSION = '2.0.5'

EXEC_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
FRAMES = '{}\\frames'.format(EXEC_PATH)  # Folder containing extracted frames
UPSCALED = '{}\\upscaled'.format(EXEC_PATH)  # Folder containing enlarges frames

# 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\"'


def process_arguments():
    """Processes CLI arguments

    This function parses all arguments
    This allows users to customize options
    for the output video.
    """
    parser = argparse.ArgumentParser()
    control_group = parser.add_argument_group('Controls')
    control_group.add_argument('--width', help='Output video width', action='store', type=int, default=False)
    control_group.add_argument('--height', help='Output video height', action='store', type=int, default=False)
    control_group.add_argument('-v', '--video', help='Specify source video file', action='store', default=False)
    control_group.add_argument('-o', '--output', help='Specify output file', action='store', default=False)
    control_group.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb')
    control_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False)
    control_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False)
    control_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False)
    return parser.parse_args()


def print_logo():
    print('__      __  _       _                  ___   __   __')
    print('\\ \\    / / (_)     | |                |__ \\  \\ \\ / /')
    print(' \\ \\  / /   _    __| |   ___    ___      ) |  \\ V /')
    print('  \\ \\/ /   | |  / _` |  / _ \\  / _ \\    / /    > <')
    print('   \\  /    | | | (_| | |  __/ | (_) |  / /_   / . \\')
    print('    \\/     |_|  \\__,_|  \\___|  \\___/  |____| /_/ \\_\\')
    print('\n               Video2X Video Enlarger')
    spaces = ((44 - len("Version " + VERSION)) // 2) * " "
    print(Avalon.FM.BD + "\n" + spaces +
          '    Version ' + VERSION + '\n' + Avalon.FM.RST)


def get_vid_info():
    """Gets original video information

    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
    """
    json_str = subprocess.check_output(
        '{} -v quiet -print_format json -show_format -show_streams \"{}\"'.format('\"' + FFMPEG_PATH + 'ffprobe.exe\"', args.video))
    return json.loads(json_str.decode('utf-8'))


def check_model_type(args):
    """
    Check if the model demanded from cli
    argument is legal.
    """
    models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
                        'anime_style_art_rgb', 'photo', 'anime_style_art_y']
    if args.model_type not in models_available:
        Avalon.error('Specified model type not found!')
        Avalon.info('Available models:')
        for model in models_available:
            print(model)
        exit(1)


def video2x():
    """Main controller for Video2X

    This function controls the flow of video conversion
    and handles all necessary functions.
    """

    check_model_type(args)

    if args.cpu:
        method = 'cpu'
    elif args.gpu:
        method = 'gpu'
    elif args.cudnn:
        method = 'cudnn'

    fm = FFMPEG('\"' + FFMPEG_PATH + 'ffmpeg.exe\"', args.output)
    w2 = WAIFU2X(WAIFU2X_PATH, method, args.model_type)

    # Clear and create directories
    if os.path.isdir(FRAMES):
        shutil.rmtree(FRAMES)
    if os.path.isdir(UPSCALED):
        shutil.rmtree(UPSCALED)
    os.mkdir(FRAMES)
    os.mkdir(UPSCALED)

    # Extract frames from video
    fm.extract_frames(args.video, FRAMES)

    info = get_vid_info()
    # Analyze original video with ffprobe and retrieve framerate
    # width, height = info['streams'][0]['width'], info['streams'][0]['height']
    framerate = float(Fraction(info['streams'][0]['avg_frame_rate']))
    Avalon.info('Framerate: {}'.format(framerate))

    # Upscale images one by one using waifu2x
    Avalon.info('Starting to upscale extracted images')
    w2.upscale(FRAMES, UPSCALED, args.width, args.height)
    Avalon.info('Conversion complete')

    # Frames to Video
    Avalon.info('Converting extracted frames into video')
    fm.convert_video(framerate, '{}x{}'.format(args.width, args.height), UPSCALED)

    # Extract and press audio in
    Avalon.info('Stripping audio track from original video')
    fm.extract_audio(args.video, UPSCALED)
    Avalon.info('Inserting audio track into new video')
    fm.insert_audio_track(UPSCALED)


# /////////////////// Execution /////////////////// #

args = process_arguments()
# Convert paths to absolute paths
args.video = os.path.abspath(args.video)
args.output = os.path.abspath(args.output)
print_logo()


# Check if FFMPEG and waifu2x are present
if not os.path.isdir(FFMPEG_PATH):
    Avalon.error('FFMPEG binaries not found')
    Avalon.error('Please specify FFMPEG binaries location in source code')
    print('Current value: {}\n'.format(FFMPEG_PATH))
    raise FileNotFoundError('FFMPEG binaries not found')
if not os.path.isfile(WAIFU2X_PATH.strip('\"')):
    Avalon.error('Waifu2x CUI executable not found')
    Avalon.error('Please specify Waifu2x CUI location in source code')
    print('Current value: {}\n'.format(WAIFU2X_PATH))
    raise FileNotFoundError('Waifu2x CUI executable not found')


# Check if arguments are valid / all necessary argument
# values are specified
if not args.video:
    Avalon.error('You need to specify the video to process')
    exit(1)
elif not args.width or not args.height:
    Avalon.error('You must specify output video width and height')
    exit(1)
elif not args.output:
    Avalon.error('You need to specify the output video name')
    exit(1)
elif not args.cpu and not args.gpu and not args.cudnn:
    Avalon.error('You need to specify the enlarging processing unit')
    exit(1)

if __name__ == '__main__':
    try:
        video2x()
    except Exception as e:
        Avalon.error('An exception occurred')
        traceback.print_exc()