Unverified Commit 75084615 authored by narugo1992's avatar narugo1992 Committed by GitHub
Browse files

Merge pull request #107 from deepghs/dev/yolo

dev(narugo): refactor yolo system design, add generic yolo modules
parents c672ade3 5280b44d
Loading
Loading
Loading
Loading
+7 −12
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 import detect_with_booru_yolo
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):
@@ -20,14 +17,12 @@ class BooruYOLODetectBenchmark(BaseBenchmark):
        self.model_name = model_name

    def load(self):
        from imgutils.detect.booru_yolo import _open_booru_yolo_model, _get_booru_yolo_labels
        _ = _open_booru_yolo_model(self.model_name)
        _ = _get_booru_yolo_labels(self.model_name)
        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.booru_yolo import _open_booru_yolo_model, _get_booru_yolo_labels
        _open_booru_yolo_model.cache_clear()
        _get_booru_yolo_labels.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)
+723 −712

File changed.

Preview size limit exceeded, changes collapsed.

+13 −15
Original line number Diff line number Diff line
import random

from benchmark import BaseBenchmark, create_plot_cli
from imgutils.detect import detect_censors
from imgutils.detect.censor import detect_censors, _REPO_ID
from imgutils.generic.yolo import _open_models_for_repo_id

_MODELS = _open_models_for_repo_id(_REPO_ID).model_names


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

    def load(self):
        from imgutils.detect.censor import _open_censor_detect_model
        _ = _open_censor_detect_model(level=self.level, version=self.version)
        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.censor import _open_censor_detect_model
        _open_censor_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_censors(image_file, level=self.level, version=self.version)
        _ = detect_censors(image_file, model_name=self.model_name)


if __name__ == '__main__':
    create_plot_cli(
        [
            ('censor v1.0 (yolov8s)', CensorDetectBenchmark('s', 'v1.0')),
            ('censor v1.0 (yolov8n)', CensorDetectBenchmark('n', 'v1.0')),
            ('censor v0.10 (yolov8s)', CensorDetectBenchmark('s', 'v0.10')),
            ('censor v0.9 (yolov8s)', CensorDetectBenchmark('s', 'v0.9')),
            # ('censor v0.8 (yolov8s)', CensorDetectBenchmark('s', 'v0.8')),
            # ('censor v0.7 (yolov8s)', CensorDetectBenchmark('s', 'v0.7')),
            (model_name, CensorDetectBenchmark(model_name))
            for model_name in _MODELS
        ],
        title='Benchmark for Anime Censor Detections',
        run_times=10,
+812 −579

File changed.

Preview size limit exceeded, changes collapsed.

+13 −13
Original line number Diff line number Diff line
import random

from benchmark import BaseBenchmark, create_plot_cli
from imgutils.detect import detect_eyes
from imgutils.detect.eye import detect_eyes, _REPO_ID
from imgutils.generic.yolo import _open_models_for_repo_id

_MODELS = _open_models_for_repo_id(_REPO_ID).model_names


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

    def load(self):
        from imgutils.detect.eye import _open_eye_detect_model
        _ = _open_eye_detect_model(level=self.level, version=self.version)
        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.eye import _open_eye_detect_model
        _open_eye_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_eyes(image_file, level=self.level, version=self.version)
        _ = detect_eyes(image_file, model_name=self.model_name)


if __name__ == '__main__':
    create_plot_cli(
        [
            ('eye v1.0 (yolov8s)', EyeDetectBenchmark('s', 'v1.0')),
            ('eye v1.0 (yolov8n)', EyeDetectBenchmark('n', 'v1.0')),
            ('eye v0.8 (yolov8s)', EyeDetectBenchmark('s', 'v0.8')),
            ('eye v0.7 (yolov8s)', EyeDetectBenchmark('s', 'v0.7')),
            (model_name, EyeDetectBenchmark(model_name))
            for model_name in _MODELS
        ],
        title='Benchmark for Anime Eyes Detections',
        run_times=10,
Loading