Spaces:
Runtime error
Runtime error
File size: 6,622 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
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
|