diff --git a/README.md b/README.md index 3fd540a..12d5808 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,8 @@ 🔥 **AnimeVideo-v3 model (动漫视频小模型)**. Please see [[*anime video models*](docs/anime_video_model.md)] and [[*comparisons*](docs/anime_comparisons.md)]
🔥 **RealESRGAN_x4plus_anime_6B** for anime images **(动漫插图模型)**. Please see [[*anime_model*](docs/anime_model.md)] -1. You can try in our website: [ARC Demo](https://arc.tencent.com/en/ai-demos/imgRestore) (now only support RealESRGAN_x4plus_anime_6B) + +1. :boom: **Add** online demo: [![Replicate](https://img.shields.io/static/v1?label=Demo&message=Replicate&color=blue)](https://replicate.com/xinntao/realesrgan). 1. [Colab Demo](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) for Real-ESRGAN **|** [Colab Demo](https://colab.research.google.com/drive/1yNl9ORUxxlL4N0keJa2SEPB61imPQd1B?usp=sharing) for Real-ESRGAN (**anime videos**) 1. Portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**. You can find more information [here](#portable-executable-files-ncnn). The ncnn implementation is in [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan) 1. You can watch enhanced animations in [Tencent Video](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html). 欢迎观看[腾讯视频动漫修复](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html) diff --git a/cog.yaml b/cog.yaml new file mode 100644 index 0000000..daa6983 --- /dev/null +++ b/cog.yaml @@ -0,0 +1,22 @@ +# This file is used for constructing replicate env +image: "r8.im/tencentarc/realesrgan" + +build: + gpu: true + python_version: "3.8" + system_packages: + - "libgl1-mesa-glx" + - "libglib2.0-0" + python_packages: + - "torch==1.7.1" + - "torchvision==0.8.2" + - "numpy==1.21.1" + - "lmdb==1.2.1" + - "opencv-python==4.5.3.56" + - "PyYAML==5.4.1" + - "tqdm==4.62.2" + - "yapf==0.31.0" + - "basicsr==1.4.2" + - "facexlib==0.2.5" + +predict: "cog_predict.py:Predictor" diff --git a/cog_predict.py b/cog_predict.py new file mode 100644 index 0000000..2be4da9 --- /dev/null +++ b/cog_predict.py @@ -0,0 +1,150 @@ +# flake8: noqa +# This file is used for deploying replicate models +# running: cog predict -i img=@inputs/00017_gray.png -i version='General - v3' -i scale=2 -i face_enhance=True -i tile=0 +# push: cog push r8.im/xinntao/realesrgan + +import os + +os.system('pip install gfpgan') +os.system('python setup.py develop') + +import cv2 +import shutil +import tempfile +import torch +from basicsr.archs.rrdbnet_arch import RRDBNet +from basicsr.archs.srvgg_arch import SRVGGNetCompact + +from realesrgan.utils import RealESRGANer + +try: + from cog import BasePredictor, Input, Path + from gfpgan import GFPGANer +except Exception: + print('please install cog and realesrgan package') + + +class Predictor(BasePredictor): + + def setup(self): + os.makedirs('output', exist_ok=True) + # download weights + if not os.path.exists('realesrgan/weights/realesr-general-x4v3.pth'): + os.system( + 'wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P ./realesrgan/weights' + ) + if not os.path.exists('realesrgan/weights/GFPGANv1.4.pth'): + os.system( + 'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P ./realesrgan/weights' + ) + if not os.path.exists('realesrgan/weights/RealESRGAN_x4plus.pth'): + os.system( + 'wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P ./realesrgan/weights' + ) + if not os.path.exists('realesrgan/weights/RealESRGAN_x4plus_anime_6B.pth'): + os.system( + 'wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P ./realesrgan/weights' + ) + if not os.path.exists('realesrgan/weights/realesr-animevideov3.pth'): + os.system( + 'wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth -P ./realesrgan/weights' + ) + + def choose_model(self, scale, version, tile=0): + half = True if torch.cuda.is_available() else False + if version == 'General - RealESRGANplus': + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4) + model_path = 'realesrgan/weights/RealESRGAN_x4plus.pth' + self.upsampler = RealESRGANer( + scale=4, model_path=model_path, model=model, tile=tile, tile_pad=10, pre_pad=0, half=half) + elif version == 'General - v3': + model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu') + model_path = 'realesrgan/weights/realesr-general-x4v3.pth' + self.upsampler = RealESRGANer( + scale=4, model_path=model_path, model=model, tile=tile, tile_pad=10, pre_pad=0, half=half) + elif version == 'Anime - anime6B': + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4) + model_path = 'realesrgan/weights/RealESRGAN_x4plus_anime_6B.pth' + self.upsampler = RealESRGANer( + scale=4, model_path=model_path, model=model, tile=tile, tile_pad=10, pre_pad=0, half=half) + elif version == 'AnimeVideo - v3': + model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu') + model_path = 'realesrgan/weights/realesr-animevideov3.pth' + self.upsampler = RealESRGANer( + scale=4, model_path=model_path, model=model, tile=tile, tile_pad=10, pre_pad=0, half=half) + + self.face_enhancer = GFPGANer( + model_path='realesrgan/weights/GFPGANv1.4.pth', + upscale=scale, + arch='clean', + channel_multiplier=2, + bg_upsampler=self.upsampler) + + def predict( + self, + img: Path = Input(description='Input'), + version: str = Input( + description='RealESRGAN version. Please see [Readme] below for more descriptions', + choices=['General - RealESRGANplus', 'General - v3', 'Anime - anime6B', 'AnimeVideo - v3'], + default='General - v3'), + scale: float = Input(description='Rescaling factor', default=2), + face_enhance: bool = Input( + description='Enhance faces with GFPGAN. Note that it does not work for anime images/vidoes', default=False), + tile: int = Input( + description= + 'Tile size. Default is 0, that is no tile. When encountering the out-of-GPU-memory issue, please specify it, e.g., 400 or 200', + default=0) + ) -> Path: + if tile <= 100 or tile is None: + tile = 0 + print(f'img: {img}. version: {version}. scale: {scale}. face_enhance: {face_enhance}. tile: {tile}.') + try: + extension = os.path.splitext(os.path.basename(str(img)))[1] + img = cv2.imread(str(img), cv2.IMREAD_UNCHANGED) + if len(img.shape) == 3 and img.shape[2] == 4: + img_mode = 'RGBA' + elif len(img.shape) == 2: + img_mode = None + img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + else: + img_mode = None + + h, w = img.shape[0:2] + if h < 300: + img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4) + + self.choose_model(scale, version, tile) + + try: + if face_enhance: + _, _, output = self.face_enhancer.enhance( + img, has_aligned=False, only_center_face=False, paste_back=True) + else: + output, _ = self.upsampler.enhance(img, outscale=scale) + except RuntimeError as error: + print('Error', error) + print('If you encounter CUDA out of memory, try to set "tile" to a smaller size, e.g., 400.') + + if img_mode == 'RGBA': # RGBA images should be saved in png format + extension = 'png' + # save_path = f'output/out.{extension}' + # cv2.imwrite(save_path, output) + out_path = Path(tempfile.mkdtemp()) / f'out.{extension}' + cv2.imwrite(str(out_path), output) + except Exception as error: + print('global exception: ', error) + finally: + clean_folder('output') + return out_path + + +def clean_folder(folder): + for filename in os.listdir(folder): + file_path = os.path.join(folder, filename) + try: + if os.path.isfile(file_path) or os.path.islink(file_path): + os.unlink(file_path) + elif os.path.isdir(file_path): + shutil.rmtree(file_path) + except Exception as e: + print(f'Failed to delete {file_path}. Reason: {e}') diff --git a/docs/Training.md b/docs/Training.md index f6494f1..77da5ea 100644 --- a/docs/Training.md +++ b/docs/Training.md @@ -166,7 +166,7 @@ You can finetune Real-ESRGAN on your own dataset. Typically, the fine-tuning pro ### Generate degraded images on the fly -Only high-resolution images are required. The low-quality images are generated with the degradation process described in Real-ESRGAN during trainig. +Only high-resolution images are required. The low-quality images are generated with the degradation process described in Real-ESRGAN during training. **1. Prepare dataset** diff --git a/inputs/00017_gray.png b/inputs/00017_gray.png new file mode 100644 index 0000000..79af68e Binary files /dev/null and b/inputs/00017_gray.png differ diff --git a/inputs/children-alpha.png b/inputs/children-alpha.png new file mode 100644 index 0000000..41dcc3b Binary files /dev/null and b/inputs/children-alpha.png differ diff --git a/requirements.txt b/requirements.txt index 8352614..0c8f3f0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ -basicsr>=1.3.3.11 -facexlib>=0.2.0.3 -gfpgan>=0.2.1 +basicsr>=1.4.2 +facexlib>=0.2.5 +gfpgan>=1.3.5 numpy opencv-python Pillow