Upload 12 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test_img/test_clip_upscaled_copy.mp4 filter=lfs diff=lfs merge=lfs -text
+test_img/test_clip_upscaled.mp4 filter=lfs diff=lfs merge=lfs -text

anime_upscaler.py

@@ -0,0 +1,125 @@
+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')

app.py

@@ -0,0 +1,7 @@
+import gradio as gr
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+iface.launch()

frame_esrgan.py

@@ -0,0 +1,72 @@
+import torch
+import torchvision
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+import cv2
+import argparse
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import image_slicer
+from image_slicer import join
+from PIL import Image
+import numpy as np
+from tqdm import tqdm
+
+def convert_from_image_to_cv2(img: Image) -> np.ndarray:
+    # return np.asarray(img)
+    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+def upscale(model_path, im_path):
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=False)
+    img = cv2.imread(im_path, cv2.IMREAD_UNCHANGED)
+    output, _ = upsampler.enhance(img, outscale=4)
+    return output
+
+def upscale_slice(model_path, image, slice):
+    width, height = Image.open(image).size
+    tiles = image_slicer.slice(image, slice, save=False)
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=False)
+    for tile in tiles:
+        output, _ = upsampler.enhance(np.array(tile.image), outscale=4)
+        tile.image = Image.fromarray(output)
+        tile.coords = (tile.coords[0]*4, tile.coords[1]*4)
+    return convert_from_image_to_cv2(join(tiles, width=width*4, height=height*4))
+
+if __name__ == '__main__':
+    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('-v', '--visualize', action='store_true', help='OPTIONAL: add this to see how image looks before and after upscale')
+    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('-r', '--resize', nargs='?', type=str, const='1920x1080', help="OPTIONAL: specify whether to resize image to a specific resolution. Specify with widthxheight, for example 1920x1080")
+    args = parser.parse_args()
+
+
+    if args.model_path and args.input and args.output:
+        if args.slice:
+            output = upscale_slice(args.model_path, args.input, args.slice)
+        else:
+            output = upscale(args.model_path, args.input)
+        if args.visualize:
+            plt.imshow(mpimg.imread(args.input))
+            plt.show()
+            plt.imshow(output)
+            plt.show()
+        if args.resize:
+            size = tuple(int(i) for i in args.resize.split('x'))
+            output = cv2.resize(output, size)
+        cv2.imwrite(args.output, output)
+    else:
+        print('Error: Missing arguments, check -h, --help for details')
+
+
+            # tiles = image_slicer.slice('tmp/{}/original/{}'.format(folder_name, i), slice, save=False)
+            # print(tiles)
+            # for tile in tiles:
+            #   up = frame_esrgan.upscale_slice(args.model_path, np.array(tile.image))
+            #   tile.image = Image.fromarray(up, 'RGB')
+            # out = join(tiles)
+            # out.save('tmp/{}/upscaled/{}'.format(folder_name, i.replace('jpg', 'png')))

readme.md

@@ -0,0 +1,43 @@
+# anime upscaler
+
+My usage of [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) to upscale anime.
+
+# Some results:
+
+## Picture
+
+Before            |  After
+:-------------------------:|:-------------------------:
+![](test_img/random_test_frame.jpg)  |  ![](test_img/random_test_frame_out.jpg)
+
+## Video
+
+### Before
+
+https://user-images.githubusercontent.com/56494763/148880280-fb6488c1-e87c-47e4-81ec-aecc3654c9d6.mp4
+
+### After
+
+https://user-images.githubusercontent.com/56494763/148880247-a99d3bd5-eb7e-4371-8de2-68d55cab801e.mp4
+
+# How To Use
+
+## Download anime model:
+
+https://github.com/xinntao/Real-ESRGAN/blob/master/docs/anime_model.md
+
+## For a single image:
+
+`python3 frame_esrgan.py -m /path/to/RealESRGAN_x4plus_anime_6B.pth -i path/to/img -o path/to/img.png`
+
+If memory usage is too high and you need to split frames and stitch back together use the `-s, --split` argument. Default number of splits is 4, but you can make as many as you want.
+
+If you want to change the final size, one can now add the `-r` or the `--resize` argument, followed by a string indicating the desired resolution, for example we would add `-r 1920x1080` to resize the output to that size.
+
+Please note that the output of a single image must be a PNG due to alpha channel. Will be fixed in the future hopefully
+
+## For a video:
+
+`python3 anime_upscaler.py -m /path/to/RealESRGAN_x4plus_anime_6B.pth -i path/to/video -o path/to/desired/output -s -a`
+
+If you do not want audio, remove the `-a` at the end

requirements.txt

@@ -0,0 +1,9 @@
+basicsr==1.4.2
+image_slicer==2.1.1
+matplotlib==3.5.1
+moviepy==1.0.3
+opencv_python==4.8.1.78
+Pillow==10.1.0
+realesrgan==0.3.0
+torchvision==0.15.2
+tqdm==4.66.1

test_img/test_clip_upscaled.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a9bdd8fe9aeef53b68add9a027e390faf92f0823dc0c26099bbca451c24ce6f
+size 15765842

test_img/test_clip_upscaled_copy.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c0c0f1875fdd41ea351f8a29b363260987916c5a8fbda0ce4bb3c86436b8f51
+size 15849639