mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-19 21:13:20 -07:00
refactoring
This commit is contained in:
Colombo
parent
246129a964
commit
8153e90ba3
13 changed files with 80 additions and 186 deletions
329
utils/DFLJPG.py
329
utils/DFLJPG.py
|
|
@ -1,329 +0,0 @@
|
|||
import pickle
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from imagelib import IEPolys
|
||||
from utils.struct_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class DFLJPG(object):
|
||||
def __init__(self):
|
||||
self.data = b""
|
||||
self.length = 0
|
||||
self.chunks = []
|
||||
self.dfl_dict = None
|
||||
self.shape = (0,0,0)
|
||||
|
||||
@staticmethod
|
||||
def load_raw(filename, loader_func=None):
|
||||
try:
|
||||
if loader_func is not None:
|
||||
data = loader_func(filename)
|
||||
else:
|
||||
with open(filename, "rb") as f:
|
||||
data = f.read()
|
||||
except:
|
||||
raise FileNotFoundError(filename)
|
||||
|
||||
try:
|
||||
inst = DFLJPG()
|
||||
inst.data = data
|
||||
inst.length = len(data)
|
||||
inst_length = inst.length
|
||||
chunks = []
|
||||
data_counter = 0
|
||||
while data_counter < inst_length:
|
||||
chunk_m_l, chunk_m_h = struct.unpack ("BB", data[data_counter:data_counter+2])
|
||||
data_counter += 2
|
||||
|
||||
if chunk_m_l != 0xFF:
|
||||
raise ValueError("No Valid JPG info")
|
||||
|
||||
chunk_name = None
|
||||
chunk_size = None
|
||||
chunk_data = None
|
||||
chunk_ex_data = None
|
||||
is_unk_chunk = False
|
||||
|
||||
if chunk_m_h & 0xF0 == 0xD0:
|
||||
n = chunk_m_h & 0x0F
|
||||
|
||||
if n >= 0 and n <= 7:
|
||||
chunk_name = "RST%d" % (n)
|
||||
chunk_size = 0
|
||||
elif n == 0x8:
|
||||
chunk_name = "SOI"
|
||||
chunk_size = 0
|
||||
if len(chunks) != 0:
|
||||
raise Exception("")
|
||||
elif n == 0x9:
|
||||
chunk_name = "EOI"
|
||||
chunk_size = 0
|
||||
elif n == 0xA:
|
||||
chunk_name = "SOS"
|
||||
elif n == 0xB:
|
||||
chunk_name = "DQT"
|
||||
elif n == 0xD:
|
||||
chunk_name = "DRI"
|
||||
chunk_size = 2
|
||||
else:
|
||||
is_unk_chunk = True
|
||||
elif chunk_m_h & 0xF0 == 0xC0:
|
||||
n = chunk_m_h & 0x0F
|
||||
if n == 0:
|
||||
chunk_name = "SOF0"
|
||||
elif n == 2:
|
||||
chunk_name = "SOF2"
|
||||
elif n == 4:
|
||||
chunk_name = "DHT"
|
||||
else:
|
||||
is_unk_chunk = True
|
||||
elif chunk_m_h & 0xF0 == 0xE0:
|
||||
n = chunk_m_h & 0x0F
|
||||
chunk_name = "APP%d" % (n)
|
||||
else:
|
||||
is_unk_chunk = True
|
||||
|
||||
if is_unk_chunk:
|
||||
raise ValueError("Unknown chunk %X" % (chunk_m_h) )
|
||||
|
||||
if chunk_size == None: #variable size
|
||||
chunk_size, = struct.unpack (">H", data[data_counter:data_counter+2])
|
||||
chunk_size -= 2
|
||||
data_counter += 2
|
||||
|
||||
if chunk_size > 0:
|
||||
chunk_data = data[data_counter:data_counter+chunk_size]
|
||||
data_counter += chunk_size
|
||||
|
||||
if chunk_name == "SOS":
|
||||
c = data_counter
|
||||
while c < inst_length and (data[c] != 0xFF or data[c+1] != 0xD9):
|
||||
c += 1
|
||||
|
||||
chunk_ex_data = data[data_counter:c]
|
||||
data_counter = c
|
||||
|
||||
chunks.append ({'name' : chunk_name,
|
||||
'm_h' : chunk_m_h,
|
||||
'data' : chunk_data,
|
||||
'ex_data' : chunk_ex_data,
|
||||
})
|
||||
inst.chunks = chunks
|
||||
|
||||
return inst
|
||||
except Exception as e:
|
||||
raise Exception ("Corrupted JPG file: %s" % (str(e)))
|
||||
|
||||
@staticmethod
|
||||
def load(filename, loader_func=None):
|
||||
try:
|
||||
inst = DFLJPG.load_raw (filename, loader_func=loader_func)
|
||||
inst.dfl_dict = None
|
||||
|
||||
for chunk in inst.chunks:
|
||||
if chunk['name'] == 'APP0':
|
||||
d, c = chunk['data'], 0
|
||||
c, id, _ = struct_unpack (d, c, "=4sB")
|
||||
|
||||
if id == b"JFIF":
|
||||
c, ver_major, ver_minor, units, Xdensity, Ydensity, Xthumbnail, Ythumbnail = struct_unpack (d, c, "=BBBHHBB")
|
||||
#if units == 0:
|
||||
# inst.shape = (Ydensity, Xdensity, 3)
|
||||
else:
|
||||
raise Exception("Unknown jpeg ID: %s" % (id) )
|
||||
elif chunk['name'] == 'SOF0' or chunk['name'] == 'SOF2':
|
||||
d, c = chunk['data'], 0
|
||||
c, precision, height, width = struct_unpack (d, c, ">BHH")
|
||||
inst.shape = (height, width, 3)
|
||||
|
||||
elif chunk['name'] == 'APP15':
|
||||
if type(chunk['data']) == bytes:
|
||||
inst.dfl_dict = pickle.loads(chunk['data'])
|
||||
|
||||
if (inst.dfl_dict is not None):
|
||||
if 'face_type' not in inst.dfl_dict:
|
||||
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
|
||||
if 'fanseg_mask' in inst.dfl_dict:
|
||||
fanseg_mask = inst.dfl_dict['fanseg_mask']
|
||||
if fanseg_mask is not None:
|
||||
numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8)
|
||||
inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
if inst.dfl_dict == None:
|
||||
return None
|
||||
|
||||
return inst
|
||||
except Exception as e:
|
||||
print (e)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def embed_dfldict(filename, dfl_dict):
|
||||
inst = DFLJPG.load_raw (filename)
|
||||
inst.setDFLDictData (dfl_dict)
|
||||
|
||||
try:
|
||||
with open(filename, "wb") as f:
|
||||
f.write ( inst.dump() )
|
||||
except:
|
||||
raise Exception( 'cannot save %s' % (filename) )
|
||||
|
||||
@staticmethod
|
||||
def embed_data(filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None,
|
||||
pitch_yaw_roll=None,
|
||||
eyebrows_expand_mod=None,
|
||||
relighted=None,
|
||||
**kwargs
|
||||
):
|
||||
|
||||
if fanseg_mask is not None:
|
||||
fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 )
|
||||
|
||||
ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] )
|
||||
|
||||
if ret and len(buf) < 60000:
|
||||
fanseg_mask = buf
|
||||
else:
|
||||
io.log_err("Unable to encode fanseg_mask for %s" % (filename) )
|
||||
fanseg_mask = None
|
||||
|
||||
DFLJPG.embed_dfldict (filename, {'face_type': face_type,
|
||||
'landmarks': landmarks,
|
||||
'ie_polys' : ie_polys.dump() if ie_polys is not None else None,
|
||||
'source_filename': source_filename,
|
||||
'source_rect': source_rect,
|
||||
'source_landmarks': source_landmarks,
|
||||
'image_to_face_mat': image_to_face_mat,
|
||||
'fanseg_mask' : fanseg_mask,
|
||||
'pitch_yaw_roll' : pitch_yaw_roll,
|
||||
'eyebrows_expand_mod' : eyebrows_expand_mod,
|
||||
'relighted' : relighted
|
||||
})
|
||||
|
||||
def embed_and_set(self, filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None,
|
||||
pitch_yaw_roll=None,
|
||||
eyebrows_expand_mod=None,
|
||||
relighted=None,
|
||||
**kwargs
|
||||
):
|
||||
if face_type is None: face_type = self.get_face_type()
|
||||
if landmarks is None: landmarks = self.get_landmarks()
|
||||
if ie_polys is None: ie_polys = self.get_ie_polys()
|
||||
if source_filename is None: source_filename = self.get_source_filename()
|
||||
if source_rect is None: source_rect = self.get_source_rect()
|
||||
if source_landmarks is None: source_landmarks = self.get_source_landmarks()
|
||||
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
|
||||
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
|
||||
if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll()
|
||||
if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod()
|
||||
if relighted is None: relighted = self.get_relighted()
|
||||
DFLJPG.embed_data (filename, face_type=face_type,
|
||||
landmarks=landmarks,
|
||||
ie_polys=ie_polys,
|
||||
source_filename=source_filename,
|
||||
source_rect=source_rect,
|
||||
source_landmarks=source_landmarks,
|
||||
image_to_face_mat=image_to_face_mat,
|
||||
fanseg_mask=fanseg_mask,
|
||||
pitch_yaw_roll=pitch_yaw_roll,
|
||||
relighted=relighted)
|
||||
|
||||
def remove_ie_polys(self):
|
||||
self.dfl_dict['ie_polys'] = None
|
||||
|
||||
def remove_fanseg_mask(self):
|
||||
self.dfl_dict['fanseg_mask'] = None
|
||||
|
||||
def remove_source_filename(self):
|
||||
self.dfl_dict['source_filename'] = None
|
||||
|
||||
def dump(self):
|
||||
data = b""
|
||||
|
||||
for chunk in self.chunks:
|
||||
data += struct.pack ("BB", 0xFF, chunk['m_h'] )
|
||||
chunk_data = chunk['data']
|
||||
if chunk_data is not None:
|
||||
data += struct.pack (">H", len(chunk_data)+2 )
|
||||
data += chunk_data
|
||||
|
||||
chunk_ex_data = chunk['ex_data']
|
||||
if chunk_ex_data is not None:
|
||||
data += chunk_ex_data
|
||||
|
||||
return data
|
||||
|
||||
def get_shape(self):
|
||||
return self.shape
|
||||
|
||||
def get_height(self):
|
||||
for chunk in self.chunks:
|
||||
if type(chunk) == IHDR:
|
||||
return chunk.height
|
||||
return 0
|
||||
|
||||
def getDFLDictData(self):
|
||||
return self.dfl_dict
|
||||
|
||||
def setDFLDictData (self, dict_data=None):
|
||||
self.dfl_dict = dict_data
|
||||
|
||||
for chunk in self.chunks:
|
||||
if chunk['name'] == 'APP15':
|
||||
self.chunks.remove(chunk)
|
||||
break
|
||||
|
||||
last_app_chunk = 0
|
||||
for i, chunk in enumerate (self.chunks):
|
||||
if chunk['m_h'] & 0xF0 == 0xE0:
|
||||
last_app_chunk = i
|
||||
|
||||
dflchunk = {'name' : 'APP15',
|
||||
'm_h' : 0xEF,
|
||||
'data' : pickle.dumps(dict_data),
|
||||
'ex_data' : None,
|
||||
}
|
||||
self.chunks.insert (last_app_chunk+1, dflchunk)
|
||||
|
||||
def get_face_type(self): return self.dfl_dict['face_type']
|
||||
def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
|
||||
def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None))
|
||||
def get_source_filename(self): return self.dfl_dict['source_filename']
|
||||
def get_source_rect(self): return self.dfl_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.dfl_dict.get ('image_to_face_mat', None)
|
||||
if mat is not None:
|
||||
return np.array (mat)
|
||||
return None
|
||||
def get_fanseg_mask(self):
|
||||
fanseg_mask = self.dfl_dict.get ('fanseg_mask', None)
|
||||
if fanseg_mask is not None:
|
||||
return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis]
|
||||
return None
|
||||
def get_pitch_yaw_roll(self):
|
||||
return self.dfl_dict.get ('pitch_yaw_roll', None)
|
||||
def get_eyebrows_expand_mod(self):
|
||||
return self.dfl_dict.get ('eyebrows_expand_mod', None)
|
||||
def get_relighted(self):
|
||||
return self.dfl_dict.get ('relighted', False)
|
||||
|
||||
439
utils/DFLPNG.py
439
utils/DFLPNG.py
|
|
@ -1,439 +0,0 @@
|
|||
import pickle
|
||||
import string
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from facelib import FaceType
|
||||
from imagelib import IEPolys
|
||||
|
||||
PNG_HEADER = b"\x89PNG\r\n\x1a\n"
|
||||
|
||||
class Chunk(object):
|
||||
def __init__(self, name=None, data=None):
|
||||
self.length = 0
|
||||
self.crc = 0
|
||||
self.name = name if name else "noNe"
|
||||
self.data = data if data else b""
|
||||
|
||||
@classmethod
|
||||
def load(cls, data):
|
||||
"""Load a chunk including header and footer"""
|
||||
inst = cls()
|
||||
if len(data) < 12:
|
||||
msg = "Chunk-data too small"
|
||||
raise ValueError(msg)
|
||||
|
||||
# chunk header & data
|
||||
(inst.length, raw_name) = struct.unpack("!I4s", data[0:8])
|
||||
inst.data = data[8:-4]
|
||||
inst.verify_length()
|
||||
inst.name = raw_name.decode("ascii")
|
||||
inst.verify_name()
|
||||
|
||||
# chunk crc
|
||||
inst.crc = struct.unpack("!I", data[8+inst.length:8+inst.length+4])[0]
|
||||
inst.verify_crc()
|
||||
|
||||
return inst
|
||||
|
||||
def dump(self, auto_crc=True, auto_length=True):
|
||||
"""Return the chunk including header and footer"""
|
||||
if auto_length: self.update_length()
|
||||
if auto_crc: self.update_crc()
|
||||
self.verify_name()
|
||||
return struct.pack("!I", self.length) + self.get_raw_name() + self.data + struct.pack("!I", self.crc)
|
||||
|
||||
def verify_length(self):
|
||||
if len(self.data) != self.length:
|
||||
msg = "Data length ({}) does not match length in chunk header ({})".format(len(self.data), self.length)
|
||||
raise ValueError(msg)
|
||||
return True
|
||||
|
||||
def verify_name(self):
|
||||
for c in self.name:
|
||||
if c not in string.ascii_letters:
|
||||
msg = "Invalid character in chunk name: {}".format(repr(self.name))
|
||||
raise ValueError(msg)
|
||||
return True
|
||||
|
||||
def verify_crc(self):
|
||||
calculated_crc = self.get_crc()
|
||||
if self.crc != calculated_crc:
|
||||
msg = "CRC mismatch: {:08X} (header), {:08X} (calculated)".format(self.crc, calculated_crc)
|
||||
raise ValueError(msg)
|
||||
return True
|
||||
|
||||
def update_length(self):
|
||||
self.length = len(self.data)
|
||||
|
||||
def update_crc(self):
|
||||
self.crc = self.get_crc()
|
||||
|
||||
def get_crc(self):
|
||||
return zlib.crc32(self.get_raw_name() + self.data)
|
||||
|
||||
def get_raw_name(self):
|
||||
return self.name if isinstance(self.name, bytes) else self.name.encode("ascii")
|
||||
|
||||
# name helper methods
|
||||
|
||||
def ancillary(self, set=None):
|
||||
"""Set and get ancillary=True/critical=False bit"""
|
||||
if set is True:
|
||||
self.name[0] = self.name[0].lower()
|
||||
elif set is False:
|
||||
self.name[0] = self.name[0].upper()
|
||||
return self.name[0].islower()
|
||||
|
||||
def private(self, set=None):
|
||||
"""Set and get private=True/public=False bit"""
|
||||
if set is True:
|
||||
self.name[1] = self.name[1].lower()
|
||||
elif set is False:
|
||||
self.name[1] = self.name[1].upper()
|
||||
return self.name[1].islower()
|
||||
|
||||
def reserved(self, set=None):
|
||||
"""Set and get reserved_valid=True/invalid=False bit"""
|
||||
if set is True:
|
||||
self.name[2] = self.name[2].upper()
|
||||
elif set is False:
|
||||
self.name[2] = self.name[2].lower()
|
||||
return self.name[2].isupper()
|
||||
|
||||
def safe_to_copy(self, set=None):
|
||||
"""Set and get save_to_copy=True/unsafe=False bit"""
|
||||
if set is True:
|
||||
self.name[3] = self.name[3].lower()
|
||||
elif set is False:
|
||||
self.name[3] = self.name[3].upper()
|
||||
return self.name[3].islower()
|
||||
|
||||
def __str__(self):
|
||||
return "<Chunk '{name}' length={length} crc={crc:08X}>".format(**self.__dict__)
|
||||
|
||||
class IHDR(Chunk):
|
||||
"""IHDR Chunk
|
||||
width, height, bit_depth, color_type, compression_method,
|
||||
filter_method, interlace_method contain the data extracted
|
||||
from the chunk. Modify those and use and build() to recreate
|
||||
the chunk. Valid values for bit_depth depend on the color_type
|
||||
and can be looked up in color_types or in the PNG specification
|
||||
|
||||
See:
|
||||
http://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html#C.IHDR
|
||||
"""
|
||||
# color types with name & allowed bit depths
|
||||
COLOR_TYPE_GRAY = 0
|
||||
COLOR_TYPE_RGB = 2
|
||||
COLOR_TYPE_PLTE = 3
|
||||
COLOR_TYPE_GRAYA = 4
|
||||
COLOR_TYPE_RGBA = 6
|
||||
color_types = {
|
||||
COLOR_TYPE_GRAY: ("Grayscale", (1,2,4,8,16)),
|
||||
COLOR_TYPE_RGB: ("RGB", (8,16)),
|
||||
COLOR_TYPE_PLTE: ("Palette", (1,2,4,8)),
|
||||
COLOR_TYPE_GRAYA: ("Greyscale+Alpha", (8,16)),
|
||||
COLOR_TYPE_RGBA: ("RGBA", (8,16)),
|
||||
}
|
||||
|
||||
def __init__(self, width=0, height=0, bit_depth=8, color_type=2, \
|
||||
compression_method=0, filter_method=0, interlace_method=0):
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.bit_depth = bit_depth
|
||||
self.color_type = color_type
|
||||
self.compression_method = compression_method
|
||||
self.filter_method = filter_method
|
||||
self.interlace_method = interlace_method
|
||||
super().__init__("IHDR")
|
||||
|
||||
@classmethod
|
||||
def load(cls, data):
|
||||
inst = super().load(data)
|
||||
fields = struct.unpack("!IIBBBBB", inst.data)
|
||||
inst.width = fields[0]
|
||||
inst.height = fields[1]
|
||||
inst.bit_depth = fields[2] # per channel
|
||||
inst.color_type = fields[3] # see specs
|
||||
inst.compression_method = fields[4] # always 0(=deflate/inflate)
|
||||
inst.filter_method = fields[5] # always 0(=adaptive filtering with 5 methods)
|
||||
inst.interlace_method = fields[6] # 0(=no interlace) or 1(=Adam7 interlace)
|
||||
return inst
|
||||
|
||||
def dump(self):
|
||||
self.data = struct.pack("!IIBBBBB", \
|
||||
self.width, self.height, self.bit_depth, self.color_type, \
|
||||
self.compression_method, self.filter_method, self.interlace_method)
|
||||
return super().dump()
|
||||
|
||||
def __str__(self):
|
||||
return "<Chunk:IHDR geometry={width}x{height} bit_depth={bit_depth} color_type={}>" \
|
||||
.format(self.color_types[self.color_type][0], **self.__dict__)
|
||||
|
||||
class IEND(Chunk):
|
||||
def __init__(self):
|
||||
super().__init__("IEND")
|
||||
|
||||
def dump(self):
|
||||
if len(self.data) != 0:
|
||||
msg = "IEND has data which is not allowed"
|
||||
raise ValueError(msg)
|
||||
if self.length != 0:
|
||||
msg = "IEND data lenght is not 0 which is not allowed"
|
||||
raise ValueError(msg)
|
||||
return super().dump()
|
||||
|
||||
def __str__(self):
|
||||
return "<Chunk:IEND>".format(**self.__dict__)
|
||||
|
||||
class DFLChunk(Chunk):
|
||||
def __init__(self, dict_data=None):
|
||||
super().__init__("fcWp")
|
||||
self.dict_data = dict_data
|
||||
|
||||
def setDictData(self, dict_data):
|
||||
self.dict_data = dict_data
|
||||
|
||||
def getDictData(self):
|
||||
return self.dict_data
|
||||
|
||||
@classmethod
|
||||
def load(cls, data):
|
||||
inst = super().load(data)
|
||||
inst.dict_data = pickle.loads( inst.data )
|
||||
return inst
|
||||
|
||||
def dump(self):
|
||||
self.data = pickle.dumps (self.dict_data)
|
||||
return super().dump()
|
||||
|
||||
chunk_map = {
|
||||
b"IHDR": IHDR,
|
||||
b"fcWp": DFLChunk,
|
||||
b"IEND": IEND
|
||||
}
|
||||
|
||||
class DFLPNG(object):
|
||||
def __init__(self):
|
||||
self.data = b""
|
||||
self.length = 0
|
||||
self.chunks = []
|
||||
self.dfl_dict = None
|
||||
|
||||
@staticmethod
|
||||
def load_raw(filename, loader_func=None):
|
||||
try:
|
||||
if loader_func is not None:
|
||||
data = loader_func(filename)
|
||||
else:
|
||||
with open(filename, "rb") as f:
|
||||
data = f.read()
|
||||
except:
|
||||
raise FileNotFoundError(filename)
|
||||
|
||||
inst = DFLPNG()
|
||||
inst.data = data
|
||||
inst.length = len(data)
|
||||
|
||||
if data[0:8] != PNG_HEADER:
|
||||
msg = "No Valid PNG header"
|
||||
raise ValueError(msg)
|
||||
|
||||
chunk_start = 8
|
||||
while chunk_start < inst.length:
|
||||
(chunk_length, chunk_name) = struct.unpack("!I4s", data[chunk_start:chunk_start+8])
|
||||
chunk_end = chunk_start + chunk_length + 12
|
||||
|
||||
chunk = chunk_map.get(chunk_name, Chunk).load(data[chunk_start:chunk_end])
|
||||
inst.chunks.append(chunk)
|
||||
chunk_start = chunk_end
|
||||
|
||||
return inst
|
||||
|
||||
@staticmethod
|
||||
def load(filename, loader_func=None):
|
||||
try:
|
||||
inst = DFLPNG.load_raw (filename, loader_func=loader_func)
|
||||
inst.dfl_dict = inst.getDFLDictData()
|
||||
|
||||
if inst.dfl_dict is not None:
|
||||
if 'face_type' not in inst.dfl_dict:
|
||||
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
|
||||
|
||||
if 'fanseg_mask' in inst.dfl_dict:
|
||||
fanseg_mask = inst.dfl_dict['fanseg_mask']
|
||||
if fanseg_mask is not None:
|
||||
numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8)
|
||||
inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
if inst.dfl_dict == None:
|
||||
return None
|
||||
|
||||
return inst
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def embed_dfldict(filename, dfl_dict):
|
||||
inst = DFLPNG.load_raw (filename)
|
||||
inst.setDFLDictData (dfl_dict)
|
||||
|
||||
try:
|
||||
with open(filename, "wb") as f:
|
||||
f.write ( inst.dump() )
|
||||
except:
|
||||
raise Exception( 'cannot save %s' % (filename) )
|
||||
|
||||
@staticmethod
|
||||
def embed_data(filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None,
|
||||
pitch_yaw_roll=None,
|
||||
eyebrows_expand_mod=None,
|
||||
relighted=None,
|
||||
**kwargs
|
||||
):
|
||||
|
||||
if fanseg_mask is not None:
|
||||
fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 )
|
||||
|
||||
ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] )
|
||||
|
||||
if ret and len(buf) < 60000:
|
||||
fanseg_mask = buf
|
||||
else:
|
||||
io.log_err("Unable to encode fanseg_mask for %s" % (filename) )
|
||||
fanseg_mask = None
|
||||
|
||||
DFLPNG.embed_dfldict (filename, {'face_type': face_type,
|
||||
'landmarks': landmarks,
|
||||
'ie_polys' : ie_polys.dump() if ie_polys is not None else None,
|
||||
'source_filename': source_filename,
|
||||
'source_rect': source_rect,
|
||||
'source_landmarks': source_landmarks,
|
||||
'image_to_face_mat':image_to_face_mat,
|
||||
'fanseg_mask' : fanseg_mask,
|
||||
'pitch_yaw_roll' : pitch_yaw_roll,
|
||||
'eyebrows_expand_mod' : eyebrows_expand_mod,
|
||||
'relighted' : relighted
|
||||
})
|
||||
|
||||
def embed_and_set(self, filename, face_type=None,
|
||||
landmarks=None,
|
||||
ie_polys=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None,
|
||||
fanseg_mask=None,
|
||||
pitch_yaw_roll=None,
|
||||
eyebrows_expand_mod=None,
|
||||
relighted=None,
|
||||
**kwargs
|
||||
):
|
||||
if face_type is None: face_type = self.get_face_type()
|
||||
if landmarks is None: landmarks = self.get_landmarks()
|
||||
if ie_polys is None: ie_polys = self.get_ie_polys()
|
||||
if source_filename is None: source_filename = self.get_source_filename()
|
||||
if source_rect is None: source_rect = self.get_source_rect()
|
||||
if source_landmarks is None: source_landmarks = self.get_source_landmarks()
|
||||
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
|
||||
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
|
||||
if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll()
|
||||
if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod()
|
||||
if relighted is None: relighted = self.get_relighted()
|
||||
|
||||
DFLPNG.embed_data (filename, face_type=face_type,
|
||||
landmarks=landmarks,
|
||||
ie_polys=ie_polys,
|
||||
source_filename=source_filename,
|
||||
source_rect=source_rect,
|
||||
source_landmarks=source_landmarks,
|
||||
image_to_face_mat=image_to_face_mat,
|
||||
fanseg_mask=fanseg_mask,
|
||||
pitch_yaw_roll=pitch_yaw_roll,
|
||||
eyebrows_expand_mod=eyebrows_expand_mod,
|
||||
relighted=relighted)
|
||||
|
||||
def remove_ie_polys(self):
|
||||
self.dfl_dict['ie_polys'] = None
|
||||
|
||||
def remove_fanseg_mask(self):
|
||||
self.dfl_dict['fanseg_mask'] = None
|
||||
|
||||
def remove_source_filename(self):
|
||||
self.dfl_dict['source_filename'] = None
|
||||
|
||||
def dump(self):
|
||||
data = PNG_HEADER
|
||||
for chunk in self.chunks:
|
||||
data += chunk.dump()
|
||||
return data
|
||||
|
||||
def get_shape(self):
|
||||
for chunk in self.chunks:
|
||||
if type(chunk) == IHDR:
|
||||
c = 3 if chunk.color_type == IHDR.COLOR_TYPE_RGB else 4
|
||||
w = chunk.width
|
||||
h = chunk.height
|
||||
return (h,w,c)
|
||||
return (0,0,0)
|
||||
|
||||
def get_height(self):
|
||||
for chunk in self.chunks:
|
||||
if type(chunk) == IHDR:
|
||||
return chunk.height
|
||||
return 0
|
||||
|
||||
def getDFLDictData(self):
|
||||
for chunk in self.chunks:
|
||||
if type(chunk) == DFLChunk:
|
||||
return chunk.getDictData()
|
||||
return None
|
||||
|
||||
def setDFLDictData (self, dict_data=None):
|
||||
for chunk in self.chunks:
|
||||
if type(chunk) == DFLChunk:
|
||||
self.chunks.remove(chunk)
|
||||
break
|
||||
|
||||
if not dict_data is None:
|
||||
chunk = DFLChunk(dict_data)
|
||||
self.chunks.insert(-1, chunk)
|
||||
|
||||
def get_face_type(self): return self.dfl_dict['face_type']
|
||||
def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
|
||||
def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None))
|
||||
def get_source_filename(self): return self.dfl_dict['source_filename']
|
||||
def get_source_rect(self): return self.dfl_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.dfl_dict.get ('image_to_face_mat', None)
|
||||
if mat is not None:
|
||||
return np.array (mat)
|
||||
return None
|
||||
def get_fanseg_mask(self):
|
||||
fanseg_mask = self.dfl_dict.get ('fanseg_mask', None)
|
||||
if fanseg_mask is not None:
|
||||
return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis]
|
||||
return None
|
||||
def get_pitch_yaw_roll(self):
|
||||
return self.dfl_dict.get ('pitch_yaw_roll', None)
|
||||
def get_eyebrows_expand_mod(self):
|
||||
return self.dfl_dict.get ('eyebrows_expand_mod', None)
|
||||
def get_relighted(self):
|
||||
return self.dfl_dict.get ('relighted', False)
|
||||
|
||||
|
||||
def __str__(self):
|
||||
return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)
|
||||
Loading…
Add table
Add a link
Reference in a new issue