Spaces:
Running
Running
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) | |