Colombo
2b7364005d
Now you can replace the head. Example: https://www.youtube.com/watch?v=xr5FHd0AdlQ Requirements: Post processing skill in Adobe After Effects or Davinci Resolve. Usage: 1) Find suitable dst footage with the monotonous background behind head 2) Use “extract head” script 3) Gather rich src headset from only one scene (same color and haircut) 4) Mask whole head for src and dst using XSeg editor 5) Train XSeg 6) Apply trained XSeg mask for src and dst headsets 7) Train SAEHD using ‘head’ face_type as regular deepfake model with DF archi. You can use pretrained model for head. Minimum recommended resolution for head is 224. 8) Extract multiple tracks, using Merger: a. Raw-rgb b. XSeg-prd mask c. XSeg-dst mask 9) Using AAE or DavinciResolve, do: a. Hide source head using XSeg-prd mask: content-aware-fill, clone-stamp, background retraction, or other technique b. Overlay new head using XSeg-dst mask Warning: Head faceset can be used for whole_face or less types of training only with XSeg masking. XSegEditor: added button ‘view trained XSeg mask’, so you can see which frames should be masked to improve mask quality.
303 lines
10 KiB
Python
303 lines
10 KiB
Python
import pickle
|
|
import struct
|
|
import traceback
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from core import imagelib
|
|
from core.imagelib import SegIEPolys
|
|
from core.interact import interact as io
|
|
from core.structex import *
|
|
from facelib import FaceType
|
|
|
|
|
|
class DFLJPG(object):
|
|
def __init__(self, filename):
|
|
self.filename = filename
|
|
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(filename)
|
|
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(f"No Valid JPG info in {filename}")
|
|
|
|
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(f"Unknown chunk {chunk_m_h} in {filename}")
|
|
# io.log_info(f"Unknown chunk {chunk_m_h} in {filename}")
|
|
|
|
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 (f"Corrupted JPG file {filename} {e}")
|
|
|
|
@staticmethod
|
|
def load(filename, loader_func=None):
|
|
try:
|
|
inst = DFLJPG.load_raw (filename, loader_func=loader_func)
|
|
inst.dfl_dict = {}
|
|
|
|
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'])
|
|
|
|
return inst
|
|
except Exception as e:
|
|
io.log_err (f'Exception occured while DFLJPG.load : {traceback.format_exc()}')
|
|
return None
|
|
|
|
def has_data(self):
|
|
return len(self.dfl_dict.keys()) != 0
|
|
|
|
def save(self):
|
|
try:
|
|
with open(self.filename, "wb") as f:
|
|
f.write ( self.dump() )
|
|
except:
|
|
raise Exception( f'cannot save {self.filename}' )
|
|
|
|
def dump(self):
|
|
data = b""
|
|
|
|
dict_data = self.dfl_dict
|
|
|
|
# Remove None keys
|
|
for key in list(dict_data.keys()):
|
|
if dict_data[key] is None:
|
|
dict_data.pop(key)
|
|
|
|
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)
|
|
|
|
|
|
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 get_dict(self):
|
|
return self.dfl_dict
|
|
|
|
def set_dict (self, dict_data=None):
|
|
self.dfl_dict = dict_data
|
|
|
|
def get_face_type(self): return self.dfl_dict.get('face_type', FaceType.toString (FaceType.FULL) )
|
|
def set_face_type(self, face_type): self.dfl_dict['face_type'] = face_type
|
|
|
|
def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
|
|
def set_landmarks(self, landmarks): self.dfl_dict['landmarks'] = landmarks
|
|
|
|
def get_eyebrows_expand_mod(self): return self.dfl_dict.get ('eyebrows_expand_mod', 1.0)
|
|
def set_eyebrows_expand_mod(self, eyebrows_expand_mod): self.dfl_dict['eyebrows_expand_mod'] = eyebrows_expand_mod
|
|
|
|
def get_source_filename(self): return self.dfl_dict.get ('source_filename', None)
|
|
def set_source_filename(self, source_filename): self.dfl_dict['source_filename'] = source_filename
|
|
|
|
def get_source_rect(self): return self.dfl_dict.get ('source_rect', None)
|
|
def set_source_rect(self, source_rect): self.dfl_dict['source_rect'] = source_rect
|
|
|
|
def get_source_landmarks(self): return np.array ( self.dfl_dict.get('source_landmarks', None) )
|
|
def set_source_landmarks(self, source_landmarks): self.dfl_dict['source_landmarks'] = 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 set_image_to_face_mat(self, image_to_face_mat): self.dfl_dict['image_to_face_mat'] = image_to_face_mat
|
|
|
|
def get_seg_ie_polys(self):
|
|
d = self.dfl_dict.get('seg_ie_polys',None)
|
|
if d is not None:
|
|
d = SegIEPolys.load(d)
|
|
else:
|
|
d = SegIEPolys()
|
|
|
|
return d
|
|
|
|
def set_seg_ie_polys(self, seg_ie_polys):
|
|
if seg_ie_polys is not None:
|
|
if not isinstance(seg_ie_polys, SegIEPolys):
|
|
raise ValueError('seg_ie_polys should be instance of SegIEPolys')
|
|
|
|
if seg_ie_polys.has_polys():
|
|
seg_ie_polys = seg_ie_polys.dump()
|
|
else:
|
|
seg_ie_polys = None
|
|
|
|
self.dfl_dict['seg_ie_polys'] = seg_ie_polys
|
|
|
|
def get_xseg_mask(self):
|
|
mask_buf = self.dfl_dict.get('xseg_mask',None)
|
|
if mask_buf is None:
|
|
return None
|
|
|
|
img = cv2.imdecode(mask_buf, cv2.IMREAD_UNCHANGED)
|
|
if len(img.shape) == 2:
|
|
img = img[...,None]
|
|
|
|
return img.astype(np.float32) / 255.0
|
|
|
|
|
|
def set_xseg_mask(self, mask_a):
|
|
if mask_a is None:
|
|
self.dfl_dict['xseg_mask'] = None
|
|
return
|
|
|
|
mask_a = imagelib.normalize_channels(mask_a, 1)
|
|
img_data = np.clip( mask_a*255, 0, 255 ).astype(np.uint8)
|
|
|
|
data_max_len = 4096
|
|
|
|
ret, buf = cv2.imencode('.png', img_data)
|
|
|
|
if not ret or len(buf) > data_max_len:
|
|
for jpeg_quality in range(100,-1,-1):
|
|
ret, buf = cv2.imencode( '.jpg', img_data, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality] )
|
|
if ret and len(buf) <= data_max_len:
|
|
break
|
|
|
|
if not ret:
|
|
raise Exception("set_xseg_mask: unable to generate image data for set_xseg_mask")
|
|
|
|
self.dfl_dict['xseg_mask'] = buf
|