Commit b3bc563e authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add docs for funcs

parent aaac99c5

imgutils/upscale/waifu2x.py

+46 −1
Original line number Diff line number Diff line 
import os.path

import os

import re

from functools import lru_cache

from typing import Optional, Mapping, Tuple
@@ -18,6 +18,12 @@ _FILENAME_PATTERN = re.compile(r'^(noise(?P<noise>\d+)_?)?(scale(?P<scale>\d+)x) 


@lru_cache()

def _load_available_version() -> Mapping[Tuple[str, str, str], Mapping[Tuple[Optional[int], int], str]]:

    """

    Load available versions of the waifu2x model from the Hugging Face Hub.



    :return: A mapping of available model versions.

    :rtype: Mapping

    """

    records = {}

    for file in _hf_fs.glob(f'{_REPOSITORY}/*/onnx_models/*/*/*.onnx'):

        segments = os.path.relpath(file, _REPOSITORY).split('/')
@@ -45,6 +51,21 @@ def _load_available_version() -> Mapping[Tuple[str, str, str], Mapping[Tuple[Opt 


@lru_cache()

def _open_waifu2x_onnx_model(version: str, model: str, type_: str, noise: Optional[int], scale: int):

    """

    Open a specific version of the waifu2x ONNX model.



    :param version: The model version.

    :type version: str

    :param model: The model type.

    :type model: str

    :param type_: The model usage type (e.g., 'art').

    :type type_: str

    :param noise: The noise level (None for no noise).

    :type noise: Optional[int]

    :param scale: The scaling factor.

    :type scale: int

    :return: The ONNX model.

    """

    _all_versions = _load_available_version()

    if (version, model, type_) in _all_versions:

        _all_k = _all_versions[(version, model, type_)]
@@ -72,6 +93,30 @@ def _single_upscale_by_waifu2x(x, version: str = '20230504', model: str = 'swin_ 
def upscale_image_by_waifu2x(image: ImageTyping, scale: int = 2, noise: Optional[int] = None,

                             version: str = '20230504', model: str = 'swin_unet', type_: str = 'art',

                             tile_size: int = 64, tile_overlap: int = 8, silent: bool = False) -> Image.Image:

    """

    Upscale an image using the waifu2x model.



    :param image: The input image.

    :type image: ImageTyping

    :param scale: The scaling factor. Default is 2.

    :type scale: int

    :param noise: The noise level. Default is None.

    :type noise: Optional[int]

    :param version: The model version. Default is '20230504'.

    :type version: str

    :param model: The model type. Default is 'swin_unet'.

    :type model: str

    :param type_: The model usage type (e.g., 'art'). Default is 'art'.

    :type type_: str

    :param tile_size: The size of processing tiles. Default is 64.

    :type tile_size: int

    :param tile_overlap: The overlap between tiles. Default is 8.

    :type tile_overlap: int

    :param silent: If True, the progress will not be displayed. Default is False.

    :type silent: bool

    :return: The upscaled image.

    :rtype: Image.Image

    """

    image = load_image(image, mode='RGB', force_background='white')

    input_ = np.array(image).astype(np.float32) / 255.0

    input_ = input_.transpose((2, 0, 1))[None, ...]