File size: 3,909 Bytes
19b3da3 9bb133c 19b3da3 9bb133c 19b3da3 9bb133c 19b3da3 a3d6c18 19b3da3 a3d6c18 19b3da3 9bb133c a3d6c18 19b3da3 a3d6c18 19b3da3 9bb133c 19b3da3 |
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 |
import math
import os
from pathlib import Path
from typing import Union
import cv2
import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from gfpgan import GFPGANer
from PIL import Image
from realesrgan import RealESRGANer
import internals.util.image as ImageUtil
from internals.util.commons import download_image
from internals.util.config import get_root_dir
class Upscaler:
__model_esrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
__model_esrgan_anime_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
__model_gfpgan_url = (
"https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
)
__loaded = False
def load(self):
if self.__loaded:
return
download_dir = Path(Path.home() / ".cache" / "realesrgan")
download_dir.mkdir(parents=True, exist_ok=True)
self.__model_path = self.__preload_model(self.__model_esrgan_url, download_dir)
self.__model_path_anime = self.__preload_model(
self.__model_esrgan_anime_url, download_dir
)
self.__model_path_gfpgan = self.__preload_model(
self.__model_gfpgan_url, download_dir
)
self.__loaded = True
def upscale(self, image: Union[str, Image.Image], resize_dimension: int) -> bytes:
model = RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=4,
)
return self.__internal_upscale(
image, resize_dimension, self.__model_path, model
)
def upscale_anime(
self, image: Union[str, Image.Image], resize_dimension: int
) -> bytes:
model = RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=4,
)
return self.__internal_upscale(
image, resize_dimension, self.__model_path_anime, model
)
def __preload_model(self, url: str, download_dir: Path):
name = url.split("/")[-1]
if not os.path.exists(str(download_dir / name)):
return load_file_from_url(
url=url,
model_dir=str(download_dir),
progress=True,
file_name=None,
)
else:
return str(download_dir / name)
def __internal_upscale(
self,
image,
resize_dimension: int,
model_path: str,
rrbdnet: RRDBNet,
) -> bytes:
if type(image) is str:
image = download_image(image)
image = ImageUtil.resize_image_to512(image)
image = ImageUtil.to_bytes(image)
if isinstance(image, Image.Image):
image = ImageUtil.to_bytes(image)
image_array = np.frombuffer(image, dtype=np.uint8)
input_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
dimension = min(input_image.shape[0], input_image.shape[1])
scale = max(math.floor(resize_dimension / dimension), 2)
os.chdir(str(Path.home() / ".cache"))
upsampler = RealESRGANer(
scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
)
face_enhancer = GFPGANer(
model_path=self.__model_path_gfpgan,
upscale=scale,
arch="clean",
channel_multiplier=2,
bg_upsampler=upsampler,
)
_, _, output = face_enhancer.enhance(
input_image, has_aligned=False, only_center_face=True, paste_back=True
)
os.chdir(get_root_dir())
out_bytes = cv2.imencode(".png", output)[1].tobytes()
return out_bytes
|