waifuc/action/split.py

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