dev(narugo): refactor another yolo model (f77496dc) · Commits · git-mirror / Imgutils

docs/source/api_doc/detect/booru_yolo_detect_benchmark.plot.py

+2 −5

Original line number	Diff line number	Diff line
		import os
		import random

		from hfutils.operate import get_hf_fs

		from benchmark import BaseBenchmark, create_plot_cli
		from imgutils.detect.booru_yolo import detect_with_booru_yolo, _REPO_ID
		from imgutils.generic.yolo import _open_models_for_repo_id

		repository = 'deepghs/booru_yolo'
		hf_fs = get_hf_fs()
		_MODELS = [
		os.path.basename(os.path.dirname(file))
		for file in hf_fs.glob(f'{repository}/*/model.onnx')
		]
		_MODELS = _open_models_for_repo_id(_REPO_ID).model_names


		class BooruYOLODetectBenchmark(BaseBenchmark):

docs/source/api_doc/detect/head_detect_benchmark.plot.py

+13 −10

Original line number	Diff line number	Diff line
		import random

		from benchmark import BaseBenchmark, create_plot_cli
		from imgutils.detect import detect_heads
		from imgutils.detect.head import detect_heads, _REPO_ID
		from imgutils.generic.yolo import _open_models_for_repo_id

		_MODELS = _open_models_for_repo_id(_REPO_ID).model_names


		class HeadDetectBenchmark(BaseBenchmark):
		def __init__(self, level):
		def __init__(self, model_name: str):
		BaseBenchmark.__init__(self)
		self.level = level
		self.model_name = model_name

		def load(self):
		from imgutils.detect.head import _open_head_detect_model
		_ = _open_head_detect_model(level=self.level)
		from imgutils.generic.yolo import _open_models_for_repo_id
		_ = _open_models_for_repo_id(_REPO_ID)._open_model(self.model_name)

		def unload(self):
		from imgutils.detect.head import _open_head_detect_model
		_open_head_detect_model.cache_clear()
		from imgutils.generic.yolo import _open_models_for_repo_id
		_open_models_for_repo_id.cache_clear()

		def run(self):
		image_file = random.choice(self.all_images)
		_ = detect_heads(image_file, level=self.level)
		_ = detect_heads(image_file, model_name=self.model_name)


		if __name__ == '__main__':
		create_plot_cli(
		[
		('head (yolov8s)', HeadDetectBenchmark('s')),
		('head (yolov8n)', HeadDetectBenchmark('n')),
		(model_name, HeadDetectBenchmark(model_name))
		for model_name in _MODELS
		],
		title='Benchmark for Anime Head Detections',
		run_times=10,

docs/source/api_doc/detect/head_detect_benchmark.plot.py.svg

deleted100644 → 0

+0 −2329

File deleted.

Preview size limit exceeded, changes collapsed.

imgutils/detect/head.py

+14 −23

Original line number	Diff line number	Diff line
		@@ -13,25 +13,15 @@ Overview:
		:align: center

		"""
		from functools import lru_cache
		from typing import List, Tuple
		from typing import List, Tuple, Optional

		from huggingface_hub import hf_hub_download
		from ..data import ImageTyping
		from ..generic import yolo_predict

		from ._yolo import _image_preprocess, _data_postprocess
		from ..data import ImageTyping, load_image, rgb_encode
		from ..utils import open_onnx_model
		_REPO_ID = 'deepghs/anime_head_detection'


		@lru_cache()
		def _open_head_detect_model(level: str = 's'):
		return open_onnx_model(hf_hub_download(
		'deepghs/imgutils-models',
		f'head_detect/head_detect_best_{level}.onnx'
		))


		def detect_heads(image: ImageTyping, level: str = 's', max_infer_size=640,
		def detect_heads(image: ImageTyping, level: str = 's', model_name: Optional[str] = None,
		conf_threshold: float = 0.3, iou_threshold: float = 0.7) \
		-> List[Tuple[Tuple[int, int, int, int], str, float]]:
		"""
		@@ -42,8 +32,8 @@ def detect_heads(image: ImageTyping, level: str = 's', max_infer_size=640,
		:param level: The model level being used can be either `s` or `n`.
		The `n` model runs faster with smaller system overhead, while the `s` model achieves higher accuracy.
		The default value is `s`.
		:param max_infer_size: The maximum image size used for model inference, if the image size exceeds this limit,
		the image will be resized and used for inference. The default value is `640` pixels.
		:param model_name: Name of the YOLO model to use, use v0 models with optional `level` when not assigned.
		:type model_name: str, optional
		:param conf_threshold: The confidence threshold, only detection results with confidence scores above
		this threshold will be returned. The default value is `0.3`.
		:param iou_threshold: The detection area coverage overlap threshold, areas with overlaps above this threshold
		@@ -68,9 +58,10 @@ def detect_heads(image: ImageTyping, level: str = 's', max_infer_size=640,
		>>> plt.imshow(detection_visualize(image, result))
		>>> plt.show()
		"""
		image = load_image(image, mode='RGB')
		new_image, old_size, new_size = _image_preprocess(image, max_infer_size)

		data = rgb_encode(new_image)[None, ...]
		output, = _open_head_detect_model(level).run(['output0'], {'images': data})
		return _data_postprocess(output[0], conf_threshold, iou_threshold, old_size, new_size, ['head'])
		return yolo_predict(
		image=image,
		repo_id=_REPO_ID,
		model_name=model_name or f'head_detect_v0_{level}',
		conf_threshold=conf_threshold,
		iou_threshold=iou_threshold,
		)

test/detect/test_head.py

+3 −2

Original line number	Diff line number	Diff line
		import pytest

		from imgutils.detect.head import _open_head_detect_model, detect_heads
		from imgutils.detect.head import detect_heads
		from imgutils.generic.yolo import _open_models_for_repo_id
		from test.testings import get_testfile


		@@ -9,7 +10,7 @@ def _release_model_after_run():
		try:
		yield
		finally:
		_open_head_detect_model.cache_clear()
		_open_models_for_repo_id.cache_clear()


		@pytest.mark.unittest