waifuc/action/filter.py

from typing import List, Optional, Literal

from imgutils.detect import detect_faces, detect_heads, detect_person
from imgutils.validate import is_monochrome, anime_classify, anime_rating

from .base import FilterAction
from ..model import ImageItem


class NoMonochromeAction(FilterAction):
    def check(self, item: ImageItem) -> bool:
        return not is_monochrome(item.image)


class OnlyMonochromeAction(FilterAction):
    def check(self, item: ImageItem) -> bool:
        return is_monochrome(item.image)


ImageClassTyping = Literal['illustration', 'bangumi', 'comic', '3d']


class ClassFilterAction(FilterAction):
    def __init__(self, classes: List[ImageClassTyping], threshold: Optional[float] = None, **kwargs):
        self.classes = classes
        self.threshold = threshold
        self.kwargs = kwargs

    def check(self, item: ImageItem) -> bool:
        cls, score = anime_classify(item.image, **self.kwargs)
        return cls in self.classes and (self.threshold is None or score >= self.threshold)


ImageRatingTyping = Literal['safe', 'r15', 'r18']


class RatingFilterAction(FilterAction):
    def __init__(self, ratings: List[ImageRatingTyping], threshold: Optional[float] = None, **kwargs):
        self.ratings = ratings
        self.threshold = threshold
        self.kwargs = kwargs

    def check(self, item: ImageItem) -> bool:
        rating, score = anime_rating(item.image, **self.kwargs)
        return rating in self.ratings and (self.threshold is None or score >= self.threshold)


class FaceCountAction(FilterAction):
    def __init__(self, count: int, level: str = 's', version: str = 'v1.4',
                 conf_threshold: float = 0.25, iou_threshold: float = 0.7):
        self.count = count
        self.level = level
        self.version = version
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

    def check(self, item: ImageItem) -> bool:
        detection = detect_faces(item.image, self.level, self.version,
                                 conf_threshold=self.conf_threshold, iou_threshold=self.iou_threshold)
        return len(detection) == self.count


class HeadCountAction(FilterAction):
    def __init__(self, count: int, level: str = 's', conf_threshold: float = 0.3, iou_threshold: float = 0.7):
        self.count = count
        self.level = level
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

    def check(self, item: ImageItem) -> bool:
        detection = detect_heads(
            item.image, self.level,
            conf_threshold=self.conf_threshold,
            iou_threshold=self.iou_threshold
        )
        return len(detection) == self.count


class PersonRatioAction(FilterAction):
    def __init__(self, ratio: float = 0.4, level: str = 'm', version: str = 'v1.1',
                 conf_threshold: float = 0.3, iou_threshold: float = 0.5):
        self.ratio = ratio
        self.level = level
        self.version = version
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

    def check(self, item: ImageItem) -> bool:
        detections = detect_person(item.image, self.level, self.version, 640, self.conf_threshold, self.iou_threshold)
        if len(detections) != 1:
            return False

        (x0, y0, x1, y1), _, _ = detections[0]
        return abs((x1 - x0) * (y1 - y0)) >= self.ratio * (item.image.width * item.image.height)


class MinSizeFilterAction(FilterAction):
    def __init__(self, min_size: int):
        self.min_size = min_size

    def check(self, item: ImageItem) -> bool:
        return min(item.image.width, item.image.height) >= self.min_size


class MinAreaFilterAction(FilterAction):
    def __init__(self, min_size: int):
        self.min_size = min_size

    def check(self, item: ImageItem) -> bool:
        return (item.image.width * item.image.height) ** 0.5 >= self.min_size