File size: 3,041 Bytes
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 |
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 PIL import Image
from realesrgan import RealESRGANer
import internals.util.image as ImageUtil
from internals.util.commons import download_image
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"
def load(self):
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
)
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)
upsampler = RealESRGANer(
scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
)
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)
output, _ = upsampler.enhance(input_image, outscale=scale)
out_bytes = cv2.imencode(".png", output)[1].tobytes()
return out_bytes
|