detection/head.py

from typing import List, Tuple

from imgutils.data import ImageTyping
from imgutils.detect import detect_heads

from .base import DeepGHSObjectDetection


class HeadDetection(DeepGHSObjectDetection):
    def __init__(self):
        DeepGHSObjectDetection.__init__(self, repo_id='deepghs/anime_head_detection')

    def _get_default_model(self) -> str:
        return 'head_detect_v0.5_s_pruned'

    def _get_default_iou_and_score(self, model_name: str) -> Tuple[float, float]:
        return 0.7, 0.3

    def _get_labels(self, model_name: str) -> List[str]:
        return ['head']

    def detect(self, image: ImageTyping, model_name: str,
               iou_threshold: float = 0.7, score_threshold: float = 0.25) -> \
            List[Tuple[Tuple[float, float, float, float], str, float]]:
        return detect_heads(image=image, model_name=model_name,
                            iou_threshold=iou_threshold, conf_threshold=score_threshold)