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 ffmpeg import FFMPEG
from fractions import Fraction
from tqdm import tqdm
from waifu2x import WAIFU2X
import argparse
import avalon_framework as avalon
import inspect
import json
import os
import shutil
import subprocess
import traceback

VERSION = '2.0.1'

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 processArguments():
    """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('-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)
    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, framerate = info['streams'][0]['width'], info['streams'][0]['height'], float(
        Fraction(info['streams'][0]['avg_frame_rate']))
    avalon.info('Framerate: {}'.format(framerate))
    final_resolution = str(width * int(args.factor)) + \
        'x' + str(height * int(args.factor))

    # Upscale images one by one using waifu2x
    avalon.info('Starting to upscale extracted images')
    for (dirpath, dirnames, filenames) in os.walk(FRAMES):
        file_list = tqdm(filenames, ascii=True)
        for file in file_list:
            if file[-4:].lower() == '.png':
                image_path = '{}\\{}'.format(dirpath, file)
                file_list.set_description('Upscaling: {}'.format(file))
                # avalon.dbgInfo('Upscaling: {}'.format(image_path))
                w2.upscale(image_path, UPSCALED, int(args.factor) *
                           width, int(args.factor) * height)
    avalon.info('Extraction complete')

    # Frames to Video
    avalon.info('Converting extracted frames into video')
    fm.to_vid(framerate, final_resolution, 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 = processArguments()
# Convert paths to absolute paths
args.video = os.path.abspath(args.video)
args.output = os.path.abspath(args.output)
print_logo()

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