Upload 88 files

waifuc/action/__init__.py

@@ -0,0 +1,13 @@
+from .align import AlignMaxSizeAction, AlignMinSizeAction, PaddingAlignAction
+from .augument import RandomFilenameAction, RandomChoiceAction, BaseRandomAction, MirrorAction
+from .background import BackgroundRemovalAction
+from .base import BaseAction, ProcessAction, FilterAction, ActionStop
+from .basic import ModeConvertAction
+from .ccip import CCIPAction
+from .count import SliceSelectAction, FirstNSelectAction
+from .filename import FileExtAction, FileOrderAction
+from .filter import NoMonochromeAction, OnlyMonochromeAction, ClassFilterAction, RatingFilterAction, FaceCountAction, \
+ HeadCountAction, PersonRatioAction, MinSizeFilterAction, MinAreaFilterAction
+from .lpips import FilterSimilarAction
+from .split import PersonSplitAction, ThreeStageSplitAction
+from .tagging import TaggingAction, TagFilterAction

waifuc/action/align.py

@@ -0,0 +1,51 @@
+from typing import Tuple
+
+from PIL import Image
+from imgutils.data import load_image
+
+from .base import ProcessAction
+from ..model import ImageItem
+
+
+class AlignMaxSizeAction(ProcessAction):
+ def __init__(self, max_size: int):
+ self._max_size = max_size
+
+ def process(self, item: ImageItem) -> ImageItem:
+ image = item.image
+ ms = max(image.width, image.height)
+ if ms > self._max_size:
+ r = ms / self._max_size
+ image = image.resize((int(image.width / r), int(image.height / r)))
+
+ return ImageItem(image, item.meta)
+
+
+class AlignMinSizeAction(ProcessAction):
+ def __init__(self, min_size: int):
+ self._min_size = min_size
+
+ def process(self, item: ImageItem) -> ImageItem:
+ image = item.image
+ ms = min(image.width, image.height)
+ if ms > self._min_size:
+ r = ms / self._min_size
+ image = image.resize((int(image.width / r), int(image.height / r)))
+
+ return ImageItem(image, item.meta)
+
+
+class PaddingAlignAction(ProcessAction):
+ def __init__(self, size: Tuple[int, int], color: str = 'white'):
+ self.width, self.height = size
+ self.color = color
+
+ def process(self, item: ImageItem) -> ImageItem:
+ image = load_image(item.image, force_background=None, mode='RGBA')
+ r = min(self.width / image.width, self.height / image.height)
+ resized = image.resize((int(image.width * r), int(image.height * r)))
+
+ new_image = Image.new('RGBA', (self.width, self.height), self.color)
+ left, top = int((new_image.width - resized.width) // 2), int((new_image.height - resized.height) // 2)
+ new_image.paste(resized, (left, top, left + resized.width, top + resized.height), resized)
+ return ImageItem(new_image.convert(item.image.mode), item.meta)

waifuc/action/augument.py

@@ -0,0 +1,68 @@
+import os.path
+import random
+from typing import Iterator, Optional, Tuple
+
+from PIL import ImageOps
+from hbutils.random import random_sha1
+
+from .base import BaseAction
+from ..model import ImageItem
+
+
+class BaseRandomAction(BaseAction):
+ def __init__(self, seed=None):
+ self.seed = seed
+ self.random = random.Random(self.seed)
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ raise NotImplementedError # pragma: no cover
+
+ def reset(self):
+ self.random = random.Random(self.seed)
+
+
+class RandomChoiceAction(BaseRandomAction):
+ def __init__(self, p=0.5, seed=None):
+ BaseRandomAction.__init__(self, seed)
+ self.p = p
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if self.random.random() <= self.p:
+ yield item
+
+
+class RandomFilenameAction(BaseRandomAction):
+ def __init__(self, ext: Optional[str] = '.png', seed=None):
+ BaseRandomAction.__init__(self, seed)
+ self.ext = ext
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if 'filename' in item.meta:
+ ext = self.ext or os.path.splitext(os.path.basename(item.meta['filename']))[0]
+ else:
+ if self.ext:
+ ext = self.ext
+ else:
+ raise NameError(f'Extension (ext) must be specified '
+ f'when filename not in metadata of image item - {item!r}.')
+
+ filename = random_sha1(rnd=self.random) + ext
+ yield ImageItem(item.image, {**item.meta, 'filename': filename})
+
+
+class MirrorAction(BaseAction):
+ def __init__(self, names: Tuple[str, str] = ('origin', 'mirror')):
+ self.origin_name, self.mirror_name = names
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if 'filename' in item.meta:
+ filebody, ext = os.path.splitext(item.meta['filename'])
+ yield ImageItem(item.image, {**item.meta, 'filename': f'{filebody}_{self.origin_name}{ext}'})
+ yield ImageItem(ImageOps.mirror(item.image),
+ {**item.meta, 'filename': f'{filebody}_{self.mirror_name}{ext}'})
+ else:
+ yield ImageItem(item.image, item.meta)
+ yield ImageItem(ImageOps.mirror(item.image), item.meta)
+
+ def reset(self):
+ pass

waifuc/action/background.py

@@ -0,0 +1,10 @@
+from imgutils.segment import segment_rgba_with_isnetis
+
+from .base import ProcessAction
+from ..model import ImageItem
+
+
+class BackgroundRemovalAction(ProcessAction):
+ def process(self, item: ImageItem) -> ImageItem:
+ _, image = segment_rgba_with_isnetis(item.image)
+ return ImageItem(image, item.meta)

waifuc/action/base.py

@@ -0,0 +1,51 @@
+from typing import Iterator, Iterable
+
+from ..model import ImageItem
+
+
+class ActionStop(Exception):
+ pass
+
+
+class BaseAction:
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ raise NotImplementedError # pragma: no cover
+
+ def iter_from(self, iter_: Iterable[ImageItem]) -> Iterator[ImageItem]:
+ for item in iter_:
+ try:
+ yield from self.iter(item)
+ except ActionStop:
+ break
+
+ def reset(self):
+ raise NotImplementedError # pragma: no cover
+
+
+class ProcessAction(BaseAction):
+ def process(self, item: ImageItem) -> ImageItem:
+ raise NotImplementedError # pragma: no cover
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ yield self.process(item)
+
+ def reset(self):
+ pass
+
+ def __call__(self, item: ImageItem) -> ImageItem:
+ return self.process(item)
+
+
+class FilterAction(BaseAction):
+ def check(self, item: ImageItem) -> bool:
+ raise NotImplementedError # pragma: no cover
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if self.check(item):
+ yield item
+
+ def reset(self):
+ pass
+
+ def __call__(self, item: ImageItem) -> bool:
+ return self.check(item)

waifuc/action/basic.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+from imgutils.data import load_image
+
+from .base import ProcessAction
+from ..model import ImageItem
+
+
+class ModeConvertAction(ProcessAction):
+ def __init__(self, mode='RGB', force_background: Optional[str] = 'white'):
+ self.mode = mode
+ self.force_background = force_background
+
+ def process(self, item: ImageItem) -> ImageItem:
+ image = load_image(item.image, mode=self.mode, force_background=self.force_background)
+ return ImageItem(image, item.meta)

waifuc/action/ccip.py

@@ -0,0 +1,151 @@
+import logging
+from enum import IntEnum
+from typing import Iterator, Optional, List, Tuple
+
+import numpy as np
+from hbutils.string import plural_word
+from hbutils.testing import disable_output
+from imgutils.metrics import ccip_extract_feature, ccip_default_threshold, ccip_clustering, ccip_batch_differences
+
+from .base import BaseAction
+from ..model import ImageItem
+
+
+class CCIPStatus(IntEnum):
+ INIT = 0x1
+ APPROACH = 0x2
+ EVAL = 0x3
+ INIT_WITH_SOURCE = 0x4
+
+
+class CCIPAction(BaseAction):
+ def __init__(self, init_source=None, *, min_val_count: int = 15, step: int = 5,
+ ratio_threshold: float = 0.6, min_clu_dump_ratio: float = 0.3, cmp_threshold: float = 0.5,
+ eps: Optional[float] = None, min_samples: Optional[int] = None,
+ model='ccip-caformer-24-randaug-pruned', threshold: Optional[float] = None):
+ self.init_source = init_source
+
+ self.min_val_count = min_val_count
+ self.step = step
+ self.ratio_threshold = ratio_threshold
+ self.min_clu_dump_ratio = min_clu_dump_ratio
+ self.cmp_threshold = cmp_threshold
+ self.eps, self.min_samples = eps, min_samples
+ self.model = model
+ self.threshold = threshold or ccip_default_threshold(self.model)
+
+ self.items = []
+ self.item_released = []
+ self.feats = []
+ if self.init_source is not None:
+ self.status = CCIPStatus.INIT_WITH_SOURCE
+ else:
+ self.status = CCIPStatus.INIT
+
+ def _extract_feature(self, item: ImageItem):
+ if 'ccip_feature' in item.meta:
+ return item.meta['ccip_feature']
+ else:
+ return ccip_extract_feature(item.image, model=self.model)
+
+ def _try_cluster(self) -> bool:
+ with disable_output():
+ clu_ids = ccip_clustering(self.feats, method='optics', model=self.model,
+ eps=self.eps, min_samples=self.min_samples)
+ clu_counts = {}
+ for id_ in clu_ids:
+ if id_ != -1:
+ clu_counts[id_] = clu_counts.get(id_, 0) + 1
+
+ clu_total = sum(clu_counts.values()) if clu_counts else 0
+ chosen_id = None
+ for id_, count in clu_counts.items():
+ if count >= clu_total * self.ratio_threshold:
+ chosen_id = id_
+ break
+
+ if chosen_id is not None:
+ feats = [feat for i, feat in enumerate(self.feats) if clu_ids[i] == chosen_id]
+ clu_dump_ratio = np.array([
+ self._compare_to_exists(feat, base_set=feats)
+ for feat in feats
+ ]).astype(float).mean()
+
+ if clu_dump_ratio >= self.min_clu_dump_ratio:
+ self.items = [item for i, item in enumerate(self.items) if clu_ids[i] == chosen_id]
+ self.item_released = [False] * len(self.items)
+ self.feats = [feat for i, feat in enumerate(self.feats) if clu_ids[i] == chosen_id]
+ return True
+ else:
+ return False
+ else:
+ return False
+
+ def _compare_to_exists(self, feat, base_set=None) -> Tuple[bool, List[int]]:
+ diffs = ccip_batch_differences([feat, *(base_set or self.feats)], model=self.model)[0, 1:]
+ matches = diffs <= self.threshold
+ return matches.astype(float).mean() >= self.cmp_threshold
+
+ def _dump_items(self) -> Iterator[ImageItem]:
+ for i in range(len(self.items)):
+ if not self.item_released[i]:
+ if self._compare_to_exists(self.feats[i]):
+ self.item_released[i] = True
+ yield self.items[i]
+
+ def _eval_iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ feat = self._extract_feature(item)
+ if self._compare_to_exists(feat):
+ self.feats.append(feat)
+ yield item
+
+ if (len(self.feats) - len(self.items)) % self.step == 0:
+ yield from self._dump_items()
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if self.status == CCIPStatus.INIT_WITH_SOURCE:
+ cnt = 0
+ logging.info('Existing anchor detected.')
+ for item_ in self.init_source:
+ self.feats.append(self._extract_feature(item_))
+ yield item_
+ cnt += 1
+ logging.info(f'{plural_word(cnt, "items")} loaded from anchor.')
+
+ self.status = CCIPStatus.EVAL
+ yield from self._eval_iter(item)
+
+ elif self.status == CCIPStatus.INIT:
+ self.items.append(item)
+ self.feats.append(self._extract_feature(item))
+
+ if len(self.items) >= self.min_val_count:
+ if self._try_cluster():
+ self.status = CCIPStatus.EVAL
+ yield from self._dump_items()
+ else:
+ self.status = CCIPStatus.APPROACH
+
+ elif self.status == CCIPStatus.APPROACH:
+ self.items.append(item)
+ self.feats.append(self._extract_feature(item))
+
+ if (len(self.items) - self.min_val_count) % self.step == 0:
+ if self._try_cluster():
+ self.status = CCIPStatus.EVAL
+ yield from self._dump_items()
+
+ elif self.status == CCIPStatus.EVAL:
+ yield from self._eval_iter(item)
+
+ else:
+ raise ValueError(f'Unknown status for {self.__class__.__name__} - {self.status!r}.')
+
+ def reset(self):
+ self.items.clear()
+ self.item_released.clear()
+ self.feats.clear()
+ if self.init_source:
+ self.status = CCIPStatus.INIT_WITH_SOURCE
+ else:
+ self.status = CCIPStatus.INIT

waifuc/action/count.py

@@ -0,0 +1,72 @@
+from typing import Iterator
+
+from .base import BaseAction, ActionStop
+from ..model import ImageItem
+
+
+class FirstNSelectAction(BaseAction):
+ def __init__(self, n: int):
+ self._n = n
+ self._passed = 0
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if self._passed < self._n:
+ yield item
+ self._passed += 1
+ else:
+ raise ActionStop
+
+ def reset(self):
+ self._passed = 0
+
+
+def _slice_process(start, stop, step):
+ start = 0 if start is None else start
+ step = 1 if step is None else step
+ if not isinstance(start, int) or start < 0:
+ raise ValueError(f'Start should be an integer no less than 0, but {start!r} found.')
+ if stop is not None and (not isinstance(stop, int) or stop < 0):
+ raise ValueError(f'Stop should be an integer no less than 0, but {stop!r} found.')
+ if not isinstance(step, int) or step < 1:
+ raise ValueError(f'Step should be an integer no less than 1, but {step!r} found.')
+
+ return start, stop, step
+
+
+class SliceSelectAction(BaseAction):
+ def __init__(self, *args):
+ if len(args) == 0:
+ slice_args = _slice_process(None, None, None)
+ elif len(args) == 1:
+ slice_args = _slice_process(None, args[0], None)
+ elif len(args) == 2:
+ slice_args = _slice_process(args[0], args[1], None)
+ elif len(args) == 3:
+ slice_args = _slice_process(args[0], args[1], args[2])
+ else:
+ raise ValueError(f'Arguments of {self.__class__.__name__} should no no more than 3, but {args!r} found.')
+
+ self._start, self._stop, self._step = slice_args
+ if self._stop is not None:
+ self._max = self._start + ((self._stop - self._start - 1) // self._step) * self._step
+ else:
+ self._max = None
+ self._current = 0
+
+ def _check_current(self):
+ if self._stop is not None and self._current >= self._stop:
+ return False
+ if self._current < self._start:
+ return False
+ return (self._current - self._start) % self._step == 0
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if self._current > self._max:
+ raise ActionStop
+ else:
+ if self._check_current():
+ yield item
+ self._current += 1
+
+ def reset(self):
+ self._current = 0

waifuc/action/filename.py

@@ -0,0 +1,46 @@
+import os
+from typing import Iterator, Optional
+
+from .base import BaseAction
+from ..model import ImageItem
+
+
+class FileExtAction(BaseAction):
+ def __init__(self, ext: str):
+ self.ext = ext
+ self.untitles = 0
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if 'filename' in item.meta:
+ filebody, _ = os.path.splitext(item.meta['filename'])
+ filename = f'{filebody}{self.ext}'
+ else:
+ self.untitles += 1
+ filename = f'untitled_{self.untitles}{self.ext}'
+
+ yield ImageItem(item.image, {**item.meta, 'filename': filename})
+
+ def reset(self):
+ self.untitles = 0
+
+
+class FileOrderAction(BaseAction):
+ def __init__(self, ext: Optional[str] = '.png'):
+ self.ext = ext
+ self._current = 0
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ self._current += 1
+ if 'filename' in item.meta:
+ _, ext = os.path.splitext(item.meta['filename'])
+ new_filename = f'{self._current}{self.ext or ext}'
+ else:
+ if not self.ext:
+ raise ValueError('No extension name provided for unnamed file.')
+ else:
+ new_filename = f'{self._current}{self.ext}'
+
+ yield ImageItem(item.image, {**item.meta, 'filename': new_filename})
+
+ def reset(self):
+ self._current = 0

waifuc/action/filter.py

@@ -0,0 +1,110 @@
+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

waifuc/action/lpips.py

@@ -0,0 +1,69 @@
+from typing import Dict, Iterator, Literal
+
+import numpy as np
+from imgutils.metrics import lpips_difference, lpips_extract_feature
+
+from .base import BaseAction
+from ..model import ImageItem
+
+
+class FeatureBucket:
+ def __init__(self, threshold: float = 0.45, capacity: int = 500, rtol=1.e-5, atol=1.e-8):
+ self.threshold = threshold
+ self.rtol, self.atol = rtol, atol
+ self.features = []
+ self.ratios = np.array([], dtype=float)
+ self.capacity = capacity
+
+ def check_duplicate(self, feat, ratio: float):
+ for id_ in np.where(np.isclose(self.ratios, ratio, rtol=self.rtol, atol=self.atol))[0]:
+ exist_feat = self.features[id_.item()]
+ if lpips_difference(exist_feat, feat) <= self.threshold:
+ return True
+
+ return False
+
+ def add(self, feat, ratio: float):
+ self.features.append(feat)
+ self.ratios = np.append(self.ratios, ratio)
+ if len(self.features) >= self.capacity * 2:
+ self.features = self.features[-self.capacity:]
+ self.ratios = self.ratios[-self.capacity:]
+
+
+FilterSimilarModeTyping = Literal['all', 'group']
+
+
+class FilterSimilarAction(BaseAction):
+ def __init__(self, mode: FilterSimilarModeTyping = 'all', threshold: float = 0.45,
+ capacity: int = 500, rtol=5.e-2, atol=2.e-2):
+ self.mode = mode
+ self.threshold, self.rtol, self.atol = threshold, rtol, atol
+ self.capacity = capacity
+ self.buckets: Dict[str, FeatureBucket] = {}
+ self.global_bucket = FeatureBucket(threshold, self.capacity, rtol, atol)
+
+ def _get_bin(self, group_id):
+ if self.mode == 'all':
+ return self.global_bucket
+ elif self.mode == 'group':
+ if group_id not in self.buckets:
+ self.buckets[group_id] = FeatureBucket(self.threshold, self.capacity, self.rtol, self.atol)
+
+ return self.buckets[group_id]
+ else:
+ raise ValueError(f'Unknown mode for filter similar action - {self.mode!r}.')
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ image = item.image
+ ratio = image.height * 1.0 / image.width
+ feat = lpips_extract_feature(image)
+ bucket = self._get_bin(item.meta.get('group_id'))
+
+ if not bucket.check_duplicate(feat, ratio):
+ bucket.add(feat, ratio)
+ yield item
+
+ def reset(self):
+ self.buckets.clear()
+ self.global_bucket = FeatureBucket(self.threshold, self.capacity, self.rtol, self.atol)

waifuc/action/split.py

@@ -0,0 +1,145 @@
+import os
+from typing import Iterator, Optional
+
+from imgutils.detect import detect_person, detect_heads, detect_halfbody, detect_eyes
+
+from .base import BaseAction
+from ..model import ImageItem
+
+
+class PersonSplitAction(BaseAction):
+ def __init__(self, keep_original: bool = False, level: str = 'm', version: str = 'v1.1',
+ conf_threshold: float = 0.3, iou_threshold: float = 0.5, keep_origin_tags: bool = False):
+ self.keep_original = keep_original
+ self.level = level
+ self.version = version
+ self.conf_threshold = conf_threshold
+ self.iou_threshold = iou_threshold
+ self.keep_origin_tags = keep_origin_tags
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ detection = detect_person(item.image, self.level, self.version,
+ conf_threshold=self.conf_threshold, iou_threshold=self.iou_threshold)
+
+ if 'filename' in item.meta:
+ filename = item.meta['filename']
+ filebody, ext = os.path.splitext(filename)
+ else:
+ filebody, ext = None, None
+
+ if self.keep_original:
+ yield item
+
+ for i, (area, type_, score) in enumerate(detection):
+ new_meta = {
+ **item.meta,
+ 'crop': {'type': type_, 'score': score},
+ }
+ if 'tags' in new_meta and not self.keep_origin_tags:
+ del new_meta['tags']
+ if filebody is not None:
+ new_meta['filename'] = f'{filebody}_person{i}{ext}'
+ yield ImageItem(item.image.crop(area), new_meta)
+
+ def reset(self):
+ pass
+
+
+class ThreeStageSplitAction(BaseAction):
+ def __init__(self, person_conf: Optional[dict] = None, halfbody_conf: Optional[dict] = None,
+ head_conf: Optional[dict] = None, head_scale: float = 1.5,
+ split_eyes: bool = False, eye_conf: Optional[dict] = None, eye_scale: float = 2.4,
+ split_person: bool = True, keep_origin_tags: bool = False):
+ self.person_conf = dict(person_conf or {})
+ self.halfbody_conf = dict(halfbody_conf or {})
+ self.head_conf = dict(head_conf or {})
+ self.eye_conf = dict(eye_conf or {})
+ self.head_scale = head_scale
+ self.eye_scale = eye_scale
+ self.split_eyes = split_eyes
+ self.split_person = split_person
+ self.keep_origin_tags = keep_origin_tags
+
+ def _split_person(self, item: ImageItem, filebody, ext):
+ if self.split_person:
+ for i, (px, type_, score) in enumerate(detect_person(item.image, **self.person_conf), start=1):
+ person_image = item.image.crop(px)
+ person_meta = {
+ **item.meta,
+ 'crop': {'type': type_, 'score': score},
+ }
+ if 'tags' in person_meta and not self.keep_origin_tags:
+ del person_meta['tags']
+ if filebody is not None:
+ person_meta['filename'] = f'{filebody}_person{i}{ext}'
+ yield i, ImageItem(person_image, person_meta)
+
+ else:
+ yield 1, item
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ if 'filename' in item.meta:
+ filename = item.meta['filename']
+ filebody, ext = os.path.splitext(filename)
+ else:
+ filebody, ext = None, None
+
+ for i, person_item in self._split_person(item, filebody, ext):
+ person_image = person_item.image
+ yield person_item
+
+ half_detects = detect_halfbody(person_image, **self.halfbody_conf)
+ if half_detects:
+ halfbody_area, halfbody_type, halfbody_score = half_detects[0]
+ halfbody_image = person_image.crop(halfbody_area)
+ halfbody_meta = {
+ **item.meta,
+ 'crop': {'type': halfbody_type, 'score': halfbody_score},
+ }
+ if 'tags' in halfbody_meta and not self.keep_origin_tags:
+ del halfbody_meta['tags']
+ if filebody is not None:
+ halfbody_meta['filename'] = f'{filebody}_person{i}_halfbody{ext}'
+ yield ImageItem(halfbody_image, halfbody_meta)
+
+ head_detects = detect_heads(person_image, **self.head_conf)
+ if head_detects:
+ (hx0, hy0, hx1, hy1), head_type, head_score = head_detects[0]
+ cx, cy = (hx0 + hx1) / 2, (hy0 + hy1) / 2
+ width, height = hx1 - hx0, hy1 - hy0
+ width = height = max(width, height) * self.head_scale
+ x0, y0 = int(max(cx - width / 2, 0)), int(max(cy - height / 2, 0))
+ x1, y1 = int(min(cx + width / 2, person_image.width)), int(min(cy + height / 2, person_image.height))
+ head_image = person_image.crop((x0, y0, x1, y1))
+ head_meta = {
+ **item.meta,
+ 'crop': {'type': head_type, 'score': head_score},
+ }
+ if 'tags' in head_meta and not self.keep_origin_tags:
+ del head_meta['tags']
+ if filebody is not None:
+ head_meta['filename'] = f'{filebody}_person{i}_head{ext}'
+ yield ImageItem(head_image, head_meta)
+
+ if self.split_eyes:
+ eye_detects = detect_eyes(head_image, **self.eye_conf)
+ for j, ((ex0, ey0, ex1, ey1), eye_type, eye_score) in enumerate(eye_detects):
+ cx, cy = (ex0 + ex1) / 2, (ey0 + ey1) / 2
+ width, height = ex1 - ex0, ey1 - ey0
+ width = height = max(width, height) * self.eye_scale
+ x0, y0 = int(max(cx - width / 2, 0)), int(max(cy - height / 2, 0))
+ x1, y1 = int(min(cx + width / 2, head_image.width)), \
+ int(min(cy + height / 2, head_image.height))
+ eye_image = head_image.crop((x0, y0, x1, y1))
+ eye_meta = {
+ **item.meta,
+ 'crop': {'type': eye_type, 'score': eye_score},
+ }
+ if 'tags' in eye_meta and not self.keep_origin_tags:
+ del eye_meta['tags']
+ if filebody is not None:
+ eye_meta['filename'] = f'{filebody}_person{i}_head_eye{j}{ext}'
+ yield ImageItem(eye_image, eye_meta)
+
+ def reset(self):
+ pass

waifuc/action/tagging.py

@@ -0,0 +1,83 @@
+from functools import partial
+from typing import Iterator, Union, List, Mapping, Literal
+
+from PIL import Image
+from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags
+
+from .base import ProcessAction, BaseAction
+from ..model import ImageItem
+
+
+def _deepdanbooru_tagging(image: Image.Image, use_real_name: bool = False,
+ general_threshold: float = 0.5, character_threshold: float = 0.5, **kwargs):
+ _ = kwargs
+ _, features, characters = get_deepdanbooru_tags(image, use_real_name, general_threshold, character_threshold)
+ return {**features, **characters}
+
+
+def _wd14_tagging(image: Image.Image, model_name: str,
+ general_threshold: float = 0.35, character_threshold: float = 0.85, **kwargs):
+ _ = kwargs
+ _, features, characters = get_wd14_tags(image, model_name, general_threshold, character_threshold)
+ return {**features, **characters}
+
+
+def _mldanbooru_tagging(image: Image.Image, use_real_name: bool = False, general_threshold: float = 0.7, **kwargs):
+ _ = kwargs
+ features = get_mldanbooru_tags(image, use_real_name, general_threshold)
+ return features
+
+
+_TAGGING_METHODS = {
+ 'deepdanbooru': _deepdanbooru_tagging,
+ 'wd14_vit': partial(_wd14_tagging, model_name='ViT'),
+ 'wd14_convnext': partial(_wd14_tagging, model_name='ConvNext'),
+ 'wd14_convnextv2': partial(_wd14_tagging, model_name='ConvNextV2'),
+ 'wd14_swinv2': partial(_wd14_tagging, model_name='SwinV2'),
+ 'mldanbooru': _mldanbooru_tagging,
+}
+
+TaggingMethodTyping = Literal[
+ 'deepdanbooru', 'wd14_vit', 'wd14_convnext', 'wd14_convnextv2', 'wd14_swinv2', 'mldanbooru']
+
+
+class TaggingAction(ProcessAction):
+ def __init__(self, method: TaggingMethodTyping = 'wd14_convnextv2', force: bool = False, **kwargs):
+ self.method = _TAGGING_METHODS[method]
+ self.force = force
+ self.kwargs = kwargs
+
+ def process(self, item: ImageItem) -> ImageItem:
+ if 'tags' in item.meta and not self.force:
+ return item
+ else:
+ tags = self.method(image=item.image, **self.kwargs)
+ return ImageItem(item.image, {**item.meta, 'tags': tags})
+
+
+class TagFilterAction(BaseAction):
+ def __init__(self, tags: Union[List[str], Mapping[str, float]],
+ method: TaggingMethodTyping = 'wd14_convnextv2', **kwargs):
+ if isinstance(tags, (list, tuple)):
+ self.tags = {tag: 1e-6 for tag in tags}
+ elif isinstance(tags, dict):
+ self.tags = dict(tags)
+ else:
+ raise TypeError(f'Unknown type of tags - {tags!r}.')
+ self.tagger = TaggingAction(method, force=False, **kwargs)
+
+ def iter(self, item: ImageItem) -> Iterator[ImageItem]:
+ item = self.tagger(item)
+ tags = item.meta['tags']
+
+ valid = True
+ for tag, min_score in self.tags.items():
+ if tags[tag] < min_score:
+ valid = False
+ break
+
+ if valid:
+ yield item
+
+ def reset(self):
+ self.tagger.reset()

waifuc/config/meta.py

@@ -0,0 +1,19 @@
+"""
+Overview:
+ Meta information for waifuc package.
+"""
+
+#: Title of this project (should be `waifuc`).
+__TITLE__ = 'waifuc'
+
+#: Version of this project.
+__VERSION__ = '0.0.1'
+
+#: Short description of the project, will be included in ``setup.py``.
+__DESCRIPTION__ = 'Efficient Train Data Collector for Anime Waifu'
+
+#: Author of this project.
+__AUTHOR__ = 'narugo1992'
+
+#: Email of the authors'.
+__AUTHOR_EMAIL__ = '[email protected]'

waifuc/export/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseExporter, SaveExporter, LocalDirectoryExporter
+from .huggingface import HuggingFaceExporter
+from .textual_inversion import TextualInversionExporter

waifuc/export/base.py

@@ -0,0 +1,79 @@
+import os.path
+from typing import Iterator
+
+from hbutils.system import remove
+from tqdm.auto import tqdm
+
+from ..model import ImageItem
+from ..utils import get_task_names
+
+
+class BaseExporter:
+ def pre_export(self):
+ raise NotImplementedError # pragma: no cover
+
+ def export_item(self, item: ImageItem):
+ raise NotImplementedError # pragma: no cover
+
+ def post_export(self):
+ raise NotImplementedError # pragma: no cover
+
+ def export_from(self, items: Iterator[ImageItem]):
+ self.pre_export()
+ names = get_task_names()
+ if names:
+ desc = f'{self.__class__.__name__} - {".".join(names)}'
+ else:
+ desc = f'{self.__class__.__name__}'
+ for item in tqdm(items, desc=desc):
+ self.export_item(item)
+ self.post_export()
+
+ def reset(self):
+ raise NotImplementedError # pragma: no cover
+
+
+class LocalDirectoryExporter(BaseExporter):
+ def __init__(self, output_dir, clear: bool = False):
+ self.output_dir = output_dir
+ self.clear = clear
+
+ def pre_export(self):
+ if self.clear and os.path.exists(self.output_dir):
+ remove(self.output_dir)
+
+ os.makedirs(self.output_dir, exist_ok=True)
+
+ def export_item(self, item: ImageItem):
+ raise NotImplementedError # pragma: no cover
+
+ def post_export(self):
+ pass
+
+ def reset(self):
+ raise NotImplementedError # pragma: no cover
+
+
+class SaveExporter(LocalDirectoryExporter):
+ def __init__(self, output_dir, clear: bool = False, no_meta: bool = False,
+ skip_when_image_exist: bool = False):
+ LocalDirectoryExporter.__init__(self, output_dir, clear)
+ self.no_meta = no_meta
+ self.untitles = 0
+ self.skip_when_image_exist = skip_when_image_exist
+
+ def export_item(self, item: ImageItem):
+ if 'filename' in item.meta:
+ filename = item.meta['filename']
+ else:
+ self.untitles += 1
+ filename = f'untited_{self.untitles}.png'
+
+ full_filename = os.path.join(self.output_dir, filename)
+ full_directory = os.path.dirname(full_filename)
+ if full_directory:
+ os.makedirs(full_directory, exist_ok=True)
+ item.save(full_filename, no_meta=self.no_meta, skip_when_image_exist=self.skip_when_image_exist)
+
+ def reset(self):
+ self.untitles = 0

waifuc/export/huggingface.py

@@ -0,0 +1,64 @@
+import os
+import zipfile
+from typing import Type, Optional, Mapping, Any
+
+from hbutils.system import TemporaryDirectory
+from huggingface_hub import HfApi
+
+from .base import LocalDirectoryExporter, BaseExporter
+from ..model import ImageItem
+
+
+class HuggingFaceExporter(BaseExporter):
+ def __init__(self, repository: str, file_in_repo: str,
+ cls: Type[LocalDirectoryExporter], args: tuple = (), kwargs: Optional[Mapping[str, Any]] = None,
+ repo_type: str = 'dataset', revision: str = 'main', hf_token: Optional[str] = None):
+ self.repository = repository
+ self.repo_type, self.revision = repo_type, revision
+ self.file_in_repo = file_in_repo
+ self.cls, self.args, self.kwargs = (cls, args, kwargs or {})
+ self._tempdir: Optional[TemporaryDirectory] = None
+ self._exporter: Optional[LocalDirectoryExporter] = None
+ self.hf_token = hf_token or os.environ.get('HF_TOKEN')
+
+ def pre_export(self):
+ self._tempdir = TemporaryDirectory()
+ self._exporter = self.cls(self._tempdir.name, *self.args, **self.kwargs)
+ self._exporter.pre_export()
+
+ def export_item(self, item: ImageItem):
+ self._exporter.export_item(item)
+
+ def post_export(self):
+ self._exporter.post_export()
+
+ # upload to huggingface
+ hf_api = HfApi(token=self.hf_token)
+ hf_api.create_repo(self.repository, repo_type=self.repo_type, exist_ok=True)
+ with TemporaryDirectory() as td:
+ zip_file = os.path.join(td, 'package.zip')
+ with zipfile.ZipFile(zip_file, mode='w') as zf:
+ for directory, _, files in os.walk(self._tempdir.name):
+ for file in files:
+ file_path = os.path.join(directory, file)
+ rel_file_path = os.path.relpath(file_path, self._tempdir.name)
+ zf.write(
+ file_path,
+ '/'.join(rel_file_path.split(os.sep))
+ )
+
+ hf_api.upload_file(
+ path_or_fileobj=zip_file,
+ repo_id=self.repository,
+ repo_type=self.repo_type,
+ path_in_repo=self.file_in_repo,
+ revision=self.revision,
+ commit_message=f'Upload {self.file_in_repo} with waifuc'
+ )
+
+ self._exporter = None
+ self._tempdir.cleanup()
+ self._tempdir = None
+
+ def reset(self):
+ pass

waifuc/export/textual_inversion.py

@@ -0,0 +1,43 @@
+import os
+
+from imgutils.tagging import tags_to_text
+
+from .base import LocalDirectoryExporter
+from ..model import ImageItem
+
+
+class TextualInversionExporter(LocalDirectoryExporter):
+ def __init__(self, output_dir: str, clear: bool = False,
+ use_spaces: bool = False, use_escape: bool = True,
+ include_score: bool = False, score_descend: bool = True,
+ skip_when_image_exist: bool = False):
+ LocalDirectoryExporter.__init__(self, output_dir, clear)
+ self.use_spaces = use_spaces
+ self.use_escape = use_escape
+ self.include_score = include_score
+ self.score_descend = score_descend
+ self.untitles = 0
+ self.skip_when_image_exist = skip_when_image_exist
+
+ def export_item(self, item: ImageItem):
+ if 'filename' in item.meta:
+ filename = item.meta['filename']
+ else:
+ self.untitles += 1
+ filename = f'untited_{self.untitles}.png'
+
+ tags = item.meta.get('tags', None) or {}
+
+ full_filename = os.path.join(self.output_dir, filename)
+ full_tagname = os.path.join(self.output_dir, os.path.splitext(filename)[0] + '.txt')
+ full_directory = os.path.dirname(full_filename)
+ if full_directory:
+ os.makedirs(full_directory, exist_ok=True)
+
+ if not self.skip_when_image_exist or not os.path.exists(full_filename):
+ item.image.save(full_filename)
+ with open(full_tagname, 'w', encoding='utf-8') as f:
+ f.write(tags_to_text(tags, self.use_spaces, self.use_escape, self.include_score, self.score_descend))
+
+ def reset(self):
+ self.untitles = 0

waifuc/model/__init__.py

@@ -0,0 +1 @@
+from .item import load_meta, dump_meta, ImageItem

waifuc/model/item.py

@@ -0,0 +1,98 @@
+import json
+import os.path
+import pickle
+from dataclasses import dataclass
+from typing import Optional
+
+from PIL import Image
+from hbutils.encoding import base64_decode, base64_encode
+from hbutils.reflection import quick_import_object
+
+NoneType = type(None)
+
+_TYPE_META = '__type'
+_BASE64_META = 'base64'
+
+
+def load_meta(data, path=()):
+ if isinstance(data, (int, float, str, NoneType)):
+ return data
+ elif isinstance(data, list):
+ return [load_meta(item, (*path, i)) for i, item in enumerate(data)]
+ elif isinstance(data, dict):
+ if _TYPE_META not in data:
+ return {key: load_meta(value, (*path, key)) for key, value in data.items()}
+ else:
+ cls, _, _ = quick_import_object(data[_TYPE_META])
+ binary = base64_decode(data[_BASE64_META])
+ obj = pickle.loads(binary)
+ if isinstance(obj, cls):
+ return obj
+ else:
+ raise TypeError(f'{cls!r} expected but {obj!r} found at {path!r}.')
+ else:
+ raise TypeError(f'Unknown type {data!r} at {path!r}.')
+
+
+def dump_meta(data, path=()):
+ if isinstance(data, (int, float, str, NoneType)):
+ return data
+ elif isinstance(data, list):
+ return [dump_meta(item, (*path, i)) for i, item in enumerate(data)]
+ elif isinstance(data, dict):
+ return {key: dump_meta(value, (*path, key)) for key, value in data.items()}
+ else:
+ cls = type(data)
+ type_str = f'{cls.__module__}.{cls.__name__}' if hasattr(cls, '__module__') else cls.__name__
+ base64_str = base64_encode(pickle.dumps(data))
+ return {
+ _TYPE_META: type_str,
+ _BASE64_META: base64_str
+ }
+
+
+@dataclass
+class ImageItem:
+ image: Image.Image
+ meta: dict
+
+ def __init__(self, image: Image.Image, meta: Optional[dict] = None):
+ self.image = image
+ self.meta = meta or {}
+
+ @classmethod
+ def _image_file_to_meta_file(cls, image_file):
+ directory, filename = os.path.split(image_file)
+ filebody, _ = os.path.splitext(filename)
+ meta_file = os.path.join(directory, f'.{filebody}_meta.json')
+ return meta_file
+
+ @classmethod
+ def load_from_image(cls, image_file):
+ image = Image.open(image_file)
+ meta_file = cls._image_file_to_meta_file(image_file)
+
+ if os.path.exists(meta_file):
+ with open(meta_file, 'r', encoding='utf-8') as f:
+ meta = load_meta(json.load(f))
+ else:
+ meta = {}
+
+ return cls(image, meta)
+
+ def save(self, image_file, no_meta: bool = False, skip_when_image_exist: bool = False):
+ if not skip_when_image_exist or not os.path.exists(image_file):
+ self.image.save(image_file)
+ if not no_meta and self.meta:
+ meta_file = self._image_file_to_meta_file(image_file)
+ with open(meta_file, 'w', encoding='utf-8') as f:
+ json.dump(dump_meta(self.meta), f)
+
+ def __repr__(self):
+ values = {'size': self.image.size}
+ for key, value in self.meta.items():
+ if isinstance(value, (int, float, str)):
+ values[key] = value
+
+ content = ', '.join(f'{key}: {values[key]!r}' for key in sorted(values.keys()))
+ return f'<{self.__class__.__name__} {content}>'

waifuc/source/__init__.py

@@ -0,0 +1,19 @@
+from .anime_pictures import AnimePicturesSource
+from .base import BaseDataSource, EmptySource
+from .compose import ParallelDataSource, ComposedDataSource
+from .danbooru import DanbooruSource, SafebooruSource, ATFBooruSource, E621LikeSource, E621Source, E926Source
+from .derpibooru import DerpibooruLikeSource, DerpibooruSource, FurbooruSource
+from .duitang import DuitangSource
+from .gchar import GcharAutoSource
+from .huashi6 import Huashi6Source
+from .konachan import KonachanLikeSource, YandeSource, KonachanSource, KonachanNetSource, LolibooruSource, \
+ Rule34LikeSource, Rule34Source, HypnoHubSource, GelbooruSource, XbooruLikeSource, XbooruSource, \
+ SafebooruOrgSource, TBIBSource
+from .local import LocalSource, LocalTISource
+from .paheal import PahealSource
+from .pixiv import BasePixivSource, PixivSearchSource, PixivUserSource, PixivRankingSource
+from .sankaku import SankakuSource, PostOrder, Rating, FileType
+from .video import VideoSource
+from .wallhaven import WallHavenSource
+from .web import WebDataSource
+from .zerochan import ZerochanSource