dev(narugo): replace for rating and teen (c90127c5) · Commits · git-mirror / Imgutils

docs/source/api_doc/validate/rating_benchmark.plot.py

+6 −5

Original line number	Diff line number	Diff line
		import random

		from benchmark import BaseBenchmark, create_plot_cli
		from imgutils.generic.classify import _open_models_for_repo_id
		from imgutils.validate import anime_rating
		from imgutils.validate.rating import _MODEL_NAMES
		from imgutils.validate.rating import _REPO_ID

		_MODEL_NAMES = _open_models_for_repo_id(_REPO_ID).model_names


		class AnimeRatingBenchmark(BaseBenchmark):
		@@ -11,12 +14,10 @@ class AnimeRatingBenchmark(BaseBenchmark):
		self.model = model

		def load(self):
		from imgutils.validate.rating import _open_anime_rating_model
		_ = _open_anime_rating_model(self.model)
		_open_models_for_repo_id(_REPO_ID)._open_model(self.model)

		def unload(self):
		from imgutils.validate.rating import _open_anime_rating_model
		_open_anime_rating_model.cache_clear()
		_open_models_for_repo_id(_REPO_ID).clear()

		def run(self):
		image_file = random.choice(self.all_images)

docs/source/api_doc/validate/rating_benchmark.plot.py.svg

deleted100644 → 0

+0 −2508

File deleted.

Preview size limit exceeded, changes collapsed.

docs/source/api_doc/validate/teen_benchmark.plot.py

+6 −5

Original line number	Diff line number	Diff line
		import random

		from benchmark import BaseBenchmark, create_plot_cli
		from imgutils.generic.classify import _open_models_for_repo_id
		from imgutils.validate import anime_teen
		from imgutils.validate.teen import _MODEL_NAMES
		from imgutils.validate.teen import _REPO_ID

		_MODEL_NAMES = _open_models_for_repo_id(_REPO_ID).model_names


		class AnimeTeenBenchmark(BaseBenchmark):
		@@ -11,12 +14,10 @@ class AnimeTeenBenchmark(BaseBenchmark):
		self.model = model

		def load(self):
		from imgutils.validate.teen import _open_anime_teen_model
		_ = _open_anime_teen_model(self.model)
		_open_models_for_repo_id(_REPO_ID)._open_model(self.model)

		def unload(self):
		from imgutils.validate.teen import _open_anime_teen_model
		_open_anime_teen_model.cache_clear()
		_open_models_for_repo_id(_REPO_ID).clear()

		def run(self):
		image_file = random.choice(self.all_images)

docs/source/api_doc/validate/teen_benchmark.plot.py.svg

deleted100644 → 0

+0 −2238

File deleted.

Preview size limit exceeded, changes collapsed.

imgutils/validate/rating.py

+7 −61

Original line number	Diff line number	Diff line
		@@ -25,68 +25,18 @@ Overview:
		it is recommended to consider using object detection-based methods**,
		such as using :func:`imgutils.detect.censor.detect_censors` to detect sensitive regions as the basis for judgment.
		"""
		import json
		from functools import lru_cache
		from typing import Tuple, Optional, Dict, List
		from typing import Tuple, Dict

		import numpy as np
		from PIL import Image
		from huggingface_hub import hf_hub_download

		from imgutils.data import rgb_encode, ImageTyping, load_image
		from imgutils.utils import open_onnx_model
		from ..data import ImageTyping
		from ..generic import classify_predict, classify_predict_score

		__all__ = [
		'anime_rating_score',
		'anime_rating',
		]

		_MODEL_NAMES = [
		'caformer_s36_plus',
		'mobilenetv3',
		'mobilenetv3_sce',
		'mobilenetv3_sce_dist',
		]
		_DEFAULT_MODEL_NAME = 'mobilenetv3_sce_dist'


		@lru_cache()
		def _open_anime_rating_model(model_name):
		return open_onnx_model(hf_hub_download(
		f'deepghs/anime_rating',
		f'{model_name}/model.onnx',
		))


		@lru_cache()
		def _open_anime_rating_labels(model_name) -> List[str]:
		with open(hf_hub_download(
		f'deepghs/anime_rating',
		f'{model_name}/meta.json',
		), 'r') as f:
		return json.load(f)['labels']


		def _img_encode(image: Image.Image, size: Tuple[int, int] = (384, 384),
		normalize: Optional[Tuple[float, float]] = (0.5, 0.5)):
		image = image.resize(size, Image.BILINEAR)
		data = rgb_encode(image, order_='CHW')

		if normalize is not None:
		mean_, std_ = normalize
		mean = np.asarray([mean_]).reshape((-1, 1, 1))
		std = np.asarray([std_]).reshape((-1, 1, 1))
		data = (data - mean) / std

		return data.astype(np.float32)


		def _raw_anime_rating(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME):
		image = load_image(image, force_background='white', mode='RGB')
		input_ = _img_encode(image)[None, ...]
		output, = _open_anime_rating_model(model_name).run(['output'], {'input': input_})

		return output
		_DEFAULT_MODEL_NAME = 'mobilenetv3_v1_pruned_ls0.1'
		_REPO_ID = 'deepghs/anime_rating'


		def anime_rating_score(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Dict[str, float]:
		@@ -127,9 +77,7 @@ def anime_rating_score(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME
		>>> anime_rating_score('rating/r18/12.jpg')
		{'safe': 6.902020231791539e-06, 'r15': 0.0005639699520543218, 'r18': 0.9994290471076965}
		"""
		output = _raw_anime_rating(image, model_name)
		values = dict(zip(_open_anime_rating_labels(model_name), map(lambda x: x.item(), output[0])))
		return values
		return classify_predict_score(image, _REPO_ID, model_name)


		def anime_rating(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Tuple[str, float]:
		@@ -170,6 +118,4 @@ def anime_rating(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> T
		>>> anime_rating('rating/r18/12.jpg')
		('r18', 0.9994290471076965)
		"""
		output = _raw_anime_rating(image, model_name)[0]
		max_id = np.argmax(output)
		return _open_anime_rating_labels(model_name)[max_id], output[max_id].item()
		return classify_predict(image, _REPO_ID, model_name)