Spaces:
Running
Running
File size: 4,666 Bytes
9df91a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import frame_esrgan
import cv2
from tqdm import tqdm
import os
import argparse
import shutil
import image_slicer
from image_slicer import join
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import subprocess
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model_path', type=str, help='REQUIRED: specify path of the model being used')
parser.add_argument('-i', '--input', type=str, help='REQUIRED: specify path of the image you want to upscale')
parser.add_argument('-o', '--output', type=str, help='REQUIRED: specify path where you want to save image')
parser.add_argument('-s', '--slice', nargs='?', type=int, const=4, help='OPTIONAL: specify weather to split frames, recommended to use to help with VRAM unless you got a fucken quadro or something' )
parser.add_argument('-a', '--audio', action='store_true', help='OPTIONAL: specify weather you want to copy audio from source as well')
parser.add_argument('-c', '--clear_temp', action='store_true', help='OPTIONAL: specify weather you want to clear temporary folder with upscaled frames after you are finished with final video')
args = parser.parse_args()
def extract_frames(vid_path, save=''):
vid = cv2.VideoCapture(vid_path)
images = []
count = 0
success, image = vid.read()
while success:
if not save:
images.append(image)
else:
print('saving frame {}...'.format(count))
cv2.imwrite(save.format(count), image)
print('done saving frame {}...'.format(count))
success, image = vid.read()
count += 1
return images
def get_fps(vid_path):
vid = cv2.VideoCapture(vid_path)
return vid.get(cv2.CAP_PROP_FPS)
def create_temp_folder(vid_path):
if os.path.exists('tmp'):
folder_name = vid_path.split('/')[-1].split('.')[0]
os.mkdir('tmp/{}'.format(folder_name))
else:
os.mkdir('tmp')
create_temp_folder(vid_path)
def get_dir(path):
if not os.path.exists(path):
os.mkdir(path)
return path
def setup_frames(vid_path, slice=None):
folder_name = vid_path.split('/')[-1].split('.')[0]
images = extract_frames(vid_path)
create_temp_folder(vid_path)
os.mkdir('tmp/{}/original'.format(folder_name))
slices = []
for i in tqdm(range(len(images))):
cv2.imwrite('tmp/{}/original'.format(folder_name)+'/frame_{}.png'.format(i), images[i])
os.mkdir('tmp/{}/upscaled'.format(folder_name))
def upscale(vid_path, slice=None):
folder_name = vid_path.split('/')[-1].split('.')[0]
print('extracting frames...')
setup_frames(vid_path)
print('upscaling...')
for i in tqdm(os.listdir('tmp/{}/original'.format(folder_name))):
if slice:
out = frame_esrgan.upscale_slice(args.model_path, 'tmp/{}/original/{}'.format(folder_name, i), slice)
else:
out = frame_esrgan.upscale(args.model_path, 'tmp/{}/original/{}'.format(folder_name, i))
cv2.imwrite('tmp/{}/upscaled/{}'.format(folder_name, i), out)
def combine_frames(video_path, new_video_path):
folder_name = video_path.split('/')[-1].split('.')[0]
images = [img for img in os.listdir('tmp/{}/upscaled'.format(folder_name))]
height, width, layers = cv2.imread('tmp/{}/upscaled/frame_0.png'.format(folder_name)).shape
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
fps = get_fps(video_path)
video = cv2.VideoWriter(new_video_path, fourcc, fps, (width, height))
for i in tqdm(range(len(images))):
video.write(cv2.imread('tmp/{}/upscaled/frame_{}.png'.format(folder_name, i)))
cv2.destroyAllWindows()
video.release()
def copy_audio(original_video_path, new_video_path, new_name=''):
#ffmpeg -i input_0.mp4 -i input_1.mp4 -c copy -map 0:v:0 -map 1:a:0 -shortest out.mp4
tmp_name = new_video_path.split('.')[0] + '_tmp.' + new_video_path.split('.')[-1]
subprocess.run([
'ffmpeg',
'-i',
new_video_path,
'-i',
original_video_path,
'-c',
'copy',
'-map',
'0:v:0',
'-map',
'1:a:0',
'-shortest',
tmp_name
])
os.replace(tmp_name, new_video_path)
if __name__ == '__main__':
if args.model_path and args.input and args.output:
try:
upscale(args.input, slice=args.slice)
combine_frames(args.input, args.output)
if args.audio:
copy_audio(args.input, args.output)
if args.clear_temp:
shutil.rmtree('tmp')
except Exception as e:
print(e)
shutil.rmtree('tmp')
|