Spaces:
Runtime error
Runtime error
File size: 3,923 Bytes
4f8ad24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
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
|