esrgan / app.py
Cropinky's picture
initial commit
896048a
raw
history blame
1.3 kB
import gradio as gr
import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from realesrgan import RealESRGANer
import gc
import os
import torch
from PIL import Image
def upscale_image(img: np.ndarray) -> np.ndarray:
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
upsampler = RealESRGANer(
scale=4,
model_path=os.path.join('weights', 'RealESRGAN_x4plus.pth'),
dni_weight=None,
model=upscale_model,
tile=0,
tile_pad=10,
pre_pad=0,
half=False,
)
output, _ = upsampler.enhance(img, outscale=4)
#cleanup
gc.collect()
torch.cuda.empty_cache()
del(upsampler)
return Image.fromarray(output)
with gr.Blocks() as demo:
#gr.Markdown("Flip text or image files using this demo.")
with gr.Tab("Upscale Image"):
with gr.Row():
image_input = gr.Image()
image_output = gr.Image()
image_button = gr.Button("Upscale")
gr.Markdown("Real esrgan implementation from [this repo](https://github.com/xinntao/Real-ESRGAN)")
image_button.click(upscale_image, inputs=image_input, outputs=image_output)
if __name__ == "__main__":
demo.launch()