DeepFaceLab/DFLIMG/DFLJPG.py

import pickle
import struct

import cv2
import numpy as np

from core.interact import interact as io
from core.structex import *
from facelib import FaceType


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(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 = 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,
                             seg_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()
        
        if seg_ie_polys is not None:
            if not isinstance(seg_ie_polys, list):
                seg_ie_polys = seg_ie_polys.dump()
                
        DFLJPG.embed_dfldict (filename, {'face_type': face_type,
                                         'landmarks': landmarks,
                                         'ie_polys' : ie_polys,
                                         'seg_ie_polys' : seg_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,
                                seg_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 seg_ie_polys is None: seg_ie_polys = self.get_seg_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,
                                     seg_ie_polys=seg_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 remove_ie_polys(self):
        self.dfl_dict['ie_polys'] = None
    
    def remove_seg_ie_polys(self):
        self.dfl_dict['seg_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_seg_ie_polys(self): return self.dfl_dict.get('seg_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)