mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-03-12 20:42:45 -07:00
326 lines
12 KiB
Python
326 lines
12 KiB
Python
import pickle
|
|
import struct
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from facelib import FaceType
|
|
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,
|
|
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
|
|
|
|
if ie_polys is not None:
|
|
if not isinstance(ie_polys, list):
|
|
ie_polys = ie_polys.dump()
|
|
|
|
DFLJPG.embed_dfldict (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,
|
|
'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,
|
|
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 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,
|
|
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 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_eyebrows_expand_mod(self):
|
|
return self.dfl_dict.get ('eyebrows_expand_mod', None)
|
|
def get_relighted(self):
|
|
return self.dfl_dict.get ('relighted', False)
|
|
|