src

3467e0f 10 days ago

4.5 kB

#!/usr/bin/env python3
"""Правила отбраковки мусорных капшенов (теги gelbooru).
Живут отдельно от обучения: `filter_dataset.py` применяет их один раз к готовому датасету,
а `train.py` о них ничего не знает — обучается на том, что осталось.
Основано на просмотре выборки: `multiple views` — это character sheet и спрайт-листы
(модель учится рисовать двух похожих персонажей рядом), `speech bubble`/`english text` —
источники каши вместо букв в кадре.
"""
import re

FILTER_RULES = {
    # листы персонажа, ракурсы, комиксы
    "views": ("multiple views", "character sheet", "concept art", "reference sheet",
              "4koma", "comic", "manga cover", "collage", "split screen", "multiple angles"),
    # речь, надписи, подписи авторов
    "text": ("speech bubble", "thought bubble", "english text", "japanese text",
             "chinese text", "korean text", "mixed-language text", "text focus",
             "signature", "artist name", "copyright name", "character name", "watermark",
             "letterboxed", "logo", "body writing", "clothes writing", "writing", "engrish"),
    # монохром и незаконченные рисунки
    "mono": ("monochrome", "greyscale", "sketch", "lineart"),
    # несколько персонажей в кадре
    "multi": ("multiple girls", "multiple boys", "2girls", "3girls", "4girls", "5girls",
              "6+girls", "2boys", "3boys", "4boys", "5boys", "6+boys", "group"),
    # низкое качество и битая анатомия: модель не должна учиться на браке
    "quality": ("lowres", "jpeg artifacts", "bad anatomy", "bad hands", "bad proportions",
                "bad quality", "worst quality", "low quality", "poorly drawn", "blurry",
                "partially visible"),
    # обложки, официальный арт, сканы, референсы
    "ref": ("official art", "cover", "promotional", "reference", "scan", "artbook", "poster"),
    # размонтированные тела и лишние конечности
    "body": ("disembodied", "amputee", "multiple heads", "extra arms", "extra legs",
             "extra limbs", "conjoined", "vore", "tentacles"),
    # обрезанные кадры: фигура не влезает в кадр
    "crop": ("out of frame", "out-of-frame", "cropped"),
    # без людей: пейзажи, предметы, еда
    "nohuman": ("no humans", "background only", "still life", "scenery", "object focus",
                "food focus"),
}

FILTER_GROUPS = tuple(FILTER_RULES) + ("long",)

# Капшены длиннее этого числа тегов — скорее всего нестандартная сцена или каша
# (медиана в датасете 16, p99 — 54): 14 % строк.
MAX_TAGS = 30

# Тег `sign` ловится только как отдельный тег: подстрокой он срабатывает на `design`
# (91 строка из 70k). Регулярка дороже сравнения подстрокой, поэтому выполняется только
# для группы text.
_SIGN_RE = re.compile(r"(^|[\s,])sign([\s,]|$)")


def filter_reason(text, groups):
    """Причина отбраковки строки или None. `text` — капшен как есть из датасета.
    Сравнение по подстроке в нижнем регистре: капшен — это список тегов через запятую,
    поэтому `"2girls"` не поймает `"12girls"` (такого тега нет), и ложных срабатываний
    на подстроках внутри других слов практически нет.
    """
    low = text.lower()
    for group in groups:
        if group == "long":  # не тег, а ограничение на число тегов
            if len(text.split(",")) > MAX_TAGS:
                return group
            continue
        for key in FILTER_RULES[group]:
            if key in low:
                return group
        if group == "text" and _SIGN_RE.search(low):
            return group
    return None

	#!/usr/bin/env python3
	"""Правила отбраковки мусорных капшенов (теги gelbooru).

	Живут отдельно от обучения: `filter_dataset.py` применяет их один раз к готовому датасету,
	а `train.py` о них ничего не знает — обучается на том, что осталось.

	Основано на просмотре выборки: `multiple views` — это character sheet и спрайт-листы
	(модель учится рисовать двух похожих персонажей рядом), `speech bubble`/`english text` —
	источники каши вместо букв в кадре.
	"""
	import re

	FILTER_RULES = {
	# листы персонажа, ракурсы, комиксы
	"views": ("multiple views", "character sheet", "concept art", "reference sheet",
	"4koma", "comic", "manga cover", "collage", "split screen", "multiple angles"),
	# речь, надписи, подписи авторов
	"text": ("speech bubble", "thought bubble", "english text", "japanese text",
	"chinese text", "korean text", "mixed-language text", "text focus",
	"signature", "artist name", "copyright name", "character name", "watermark",
	"letterboxed", "logo", "body writing", "clothes writing", "writing", "engrish"),
	# монохром и незаконченные рисунки
	"mono": ("monochrome", "greyscale", "sketch", "lineart"),
	# несколько персонажей в кадре
	"multi": ("multiple girls", "multiple boys", "2girls", "3girls", "4girls", "5girls",
	"6+girls", "2boys", "3boys", "4boys", "5boys", "6+boys", "group"),
	# низкое качество и битая анатомия: модель не должна учиться на браке
	"quality": ("lowres", "jpeg artifacts", "bad anatomy", "bad hands", "bad proportions",
	"bad quality", "worst quality", "low quality", "poorly drawn", "blurry",
	"partially visible"),
	# обложки, официальный арт, сканы, референсы
	"ref": ("official art", "cover", "promotional", "reference", "scan", "artbook", "poster"),
	# размонтированные тела и лишние конечности
	"body": ("disembodied", "amputee", "multiple heads", "extra arms", "extra legs",
	"extra limbs", "conjoined", "vore", "tentacles"),
	# обрезанные кадры: фигура не влезает в кадр
	"crop": ("out of frame", "out-of-frame", "cropped"),
	# без людей: пейзажи, предметы, еда
	"nohuman": ("no humans", "background only", "still life", "scenery", "object focus",
	"food focus"),
	}

	FILTER_GROUPS = tuple(FILTER_RULES) + ("long",)

	# Капшены длиннее этого числа тегов — скорее всего нестандартная сцена или каша
	# (медиана в датасете 16, p99 — 54): 14 % строк.
	MAX_TAGS = 30

	# Тег `sign` ловится только как отдельный тег: подстрокой он срабатывает на `design`
	# (91 строка из 70k). Регулярка дороже сравнения подстрокой, поэтому выполняется только
	# для группы text.
	_SIGN_RE = re.compile(r"(^\|[\s,])sign([\s,]\|$)")


	def filter_reason(text, groups):
	"""Причина отбраковки строки или None. `text` — капшен как есть из датасета.

	Сравнение по подстроке в нижнем регистре: капшен — это список тегов через запятую,
	поэтому `"2girls"` не поймает `"12girls"` (такого тега нет), и ложных срабатываний
	на подстроках внутри других слов практически нет.
	"""
	low = text.lower()
	for group in groups:
	if group == "long": # не тег, а ограничение на число тегов
	if len(text.split(",")) > MAX_TAGS:
	return group
	continue
	for key in FILTER_RULES[group]:
	if key in low:
	return group
	if group == "text" and _SIGN_RE.search(low):
	return group
	return None