Commit 66ef53b1 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): try add aesthetic score

parent 47d39c71

docs/source/api_doc/metrics/aesthetic.benchmark.py

0 → 100644
+29 −0
Original line number Diff line number Diff line 
import random



from benchmark import BaseBenchmark, create_plot_cli

from imgutils.metrics import get_aesthetic_score





class AestheticBenchmark(BaseBenchmark):

    def load(self):

        from imgutils.metrics.aesthetic import _open_aesthetic_model

        _ = _open_aesthetic_model()



    def unload(self):

        from imgutils.metrics.aesthetic import _open_aesthetic_model

        _open_aesthetic_model.cache_clear()



    def run(self):

        image_file = random.choice(self.all_images)

        _ = get_aesthetic_score(image_file)





if __name__ == '__main__':

    create_plot_cli(

        [

            ('aesthetic', AestheticBenchmark()),

        ],

        title='Benchmark for Aesthetic Models',

        run_times=10,

        try_times=20,

    )()

docs/source/api_doc/metrics/aesthetic.rst

0 → 100644
+15 −0
Original line number Diff line number Diff line 
imgutils.metrics.aesthetic

=================================



.. currentmodule:: imgutils.metrics.aesthetic



.. automodule:: imgutils.metrics.aesthetic





get_aesthetic_score

--------------------------------------



.. autofunction:: get_aesthetic_score

docs/source/api_doc/metrics/index.rst

+1 −0
Original line number Diff line number Diff line
@@ -9,5 +9,6 @@ imgutils.metrics 
.. toctree::

    :maxdepth: 3



    aesthetic

    lpips

    psnr_

imgutils/metrics/__init__.py

+1 −0
Original line number Diff line number Diff line
@@ -2,5 +2,6 @@ 
Overview:

    Tools for computing visual metrics.

"""

from .aesthetic import *

from .lpips import *

from .psnr_ import *

imgutils/metrics/aesthetic.py

0 → 100644
+50 −0
Original line number Diff line number Diff line 
"""

Overview:

    A tool for measuring the aesthetic level of anime images, with the model

    obtained from ` <https://huggingface.co/skytnt/anime-aesthetic>`_.



    This is an overall benchmark of all the operations in LPIPS models:



    .. image:: aesthetic.benchmark.py.svg

        :align: center

"""

from functools import lru_cache



import cv2

import numpy as np

from PIL import Image

from huggingface_hub import hf_hub_download



from ..data import ImageTyping, load_image

from ..utils import open_onnx_model



__all__ = [

    'get_aesthetic_score',

]





@lru_cache()

def _open_aesthetic_model():

    return open_onnx_model(hf_hub_download(

        repo_id="skytnt/anime-aesthetic",

        filename="model.onnx"

    ))





def _preprocess(image: Image.Image):

    assert image.mode == 'RGB'

    img = np.array(image).astype(np.float32) / 255

    s = 768

    h, w = img.shape[:-1]

    h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)

    ph, pw = s - h, s - w

    img_input = np.zeros([s, s, 3], dtype=np.float32)

    img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h))

    img_input = np.transpose(img_input, (2, 0, 1))

    return img_input[np.newaxis, :]





def get_aesthetic_score(image: ImageTyping):

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

    retval, *_ = _open_aesthetic_model().run(None, {'img': _preprocess(image)})

    return float(retval.item())