File size: 4,745 Bytes
fcd5579 |
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 |
from enum import IntEnum
from pathlib import Path
import cv2
import numpy as np
from core.cv2ex import *
from facelib import LandmarksProcessor
from core import imagelib
from core.imagelib import SegIEPolys
class SampleType(IntEnum):
IMAGE = 0 #raw image
FACE_BEGIN = 1
FACE = 1 #aligned face unsorted
FACE_PERSON = 2 #aligned face person
FACE_TEMPORAL_SORTED = 3 #sorted by source filename
FACE_END = 3
QTY = 4
class Sample(object):
__slots__ = ['sample_type',
'filename',
'face_type',
'shape',
'landmarks',
'seg_ie_polys',
'xseg_mask',
'xseg_mask_compressed',
'eyebrows_expand_mod',
'source_filename',
'person_name',
'pitch_yaw_roll',
'_filename_offset_size',
]
def __init__(self, sample_type=None,
filename=None,
face_type=None,
shape=None,
landmarks=None,
seg_ie_polys=None,
xseg_mask=None,
xseg_mask_compressed=None,
eyebrows_expand_mod=None,
source_filename=None,
person_name=None,
pitch_yaw_roll=None,
**kwargs):
self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
self.filename = filename
self.face_type = face_type
self.shape = shape
self.landmarks = np.array(landmarks) if landmarks is not None else None
if isinstance(seg_ie_polys, SegIEPolys):
self.seg_ie_polys = seg_ie_polys
else:
self.seg_ie_polys = SegIEPolys.load(seg_ie_polys)
self.xseg_mask = xseg_mask
self.xseg_mask_compressed = xseg_mask_compressed
if self.xseg_mask_compressed is None and self.xseg_mask is not None:
xseg_mask = np.clip( imagelib.normalize_channels(xseg_mask, 1)*255, 0, 255 ).astype(np.uint8)
ret, xseg_mask_compressed = cv2.imencode('.png', xseg_mask)
if not ret:
raise Exception("Sample(): unable to generate xseg_mask_compressed")
self.xseg_mask_compressed = xseg_mask_compressed
self.xseg_mask = None
self.eyebrows_expand_mod = eyebrows_expand_mod if eyebrows_expand_mod is not None else 1.0
self.source_filename = source_filename
self.person_name = person_name
self.pitch_yaw_roll = pitch_yaw_roll
self._filename_offset_size = None
def has_xseg_mask(self):
return self.xseg_mask is not None or self.xseg_mask_compressed is not None
def get_xseg_mask(self):
if self.xseg_mask_compressed is not None:
xseg_mask = cv2.imdecode(self.xseg_mask_compressed, cv2.IMREAD_UNCHANGED)
if len(xseg_mask.shape) == 2:
xseg_mask = xseg_mask[...,None]
return xseg_mask.astype(np.float32) / 255.0
return self.xseg_mask
def get_pitch_yaw_roll(self):
if self.pitch_yaw_roll is None:
self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(self.landmarks, size=self.shape[1])
return self.pitch_yaw_roll
def set_filename_offset_size(self, filename, offset, size):
self._filename_offset_size = (filename, offset, size)
def read_raw_file(self, filename=None):
if self._filename_offset_size is not None:
filename, offset, size = self._filename_offset_size
with open(filename, "rb") as f:
f.seek( offset, 0)
return f.read (size)
else:
with open(filename, "rb") as f:
return f.read()
def load_bgr(self):
img = cv2_imread (self.filename, loader_func=self.read_raw_file).astype(np.float32) / 255.0
return img
def get_config(self):
return {'sample_type': self.sample_type,
'filename': self.filename,
'face_type': self.face_type,
'shape': self.shape,
'landmarks': self.landmarks.tolist(),
'seg_ie_polys': self.seg_ie_polys.dump(),
'xseg_mask' : self.xseg_mask,
'xseg_mask_compressed' : self.xseg_mask_compressed,
'eyebrows_expand_mod': self.eyebrows_expand_mod,
'source_filename': self.source_filename,
'person_name': self.person_name
}
|