diff --git a/main.py b/main.py index dace3c7..ce878fe 100644 --- a/main.py +++ b/main.py @@ -53,8 +53,6 @@ if __name__ == "__main__": extract_parser.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.") extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") extract_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") - - extract_parser.set_defaults (func=process_extract) def process_sort(arguments): @@ -66,6 +64,17 @@ if __name__ == "__main__": sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "final", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) sort_parser.set_defaults (func=process_sort) + def process_util(arguments): + from mainscripts import Util + + if arguments.convert_png_to_jpg: + Util.convert_png_to_jpg_folder (input_path=arguments.input_dir) + + util_parser = subparsers.add_parser( "util", help="Utilities.") + util_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") + util_parser.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") + util_parser.set_defaults (func=process_util) + def process_train(arguments): from mainscripts import Trainer Trainer.main ( diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 74d37c1..47500b0 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -6,6 +6,7 @@ from utils import Path_utils import cv2 from tqdm import tqdm from utils.DFLPNG import DFLPNG +from utils.DFLJPG import DFLJPG from utils import image_utils import shutil import numpy as np @@ -163,7 +164,14 @@ class ConvertSubprocessor(SubprocessorBase): cv2.waitKey(0) faces_processed = 1 elif self.converter.get_mode() == ConverterBase.MODE_IMAGE_WITH_LANDMARKS: - image_landmarks = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True ).get_landmarks() + if filename_path.suffix == '.png': + dflimg = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True ) + elif filename_path.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filename_path), throw_on_no_embedded_data=True ) + else: + raise Exception ("%s is not a dfl image file" % (filename_path.name) ) + + image_landmarks = dflimg.get_landmarks() image = self.converter.convert_image(image, image_landmarks, self.debug) if self.debug: @@ -259,16 +267,21 @@ def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_o alignments = {} aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) - for filename in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ): - dflpng = DFLPNG.load( str(filename), print_on_no_embedded_data=True ) - if dflpng is None: - continue + for filepath in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ): + filepath = Path(filepath) - source_filename_stem = Path( dflpng.get_source_filename() ).stem + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) + + source_filename_stem = Path( dflimg.get_source_filename() ).stem if source_filename_stem not in alignments.keys(): alignments[ source_filename_stem ] = [] - alignments[ source_filename_stem ].append (dflpng.get_source_landmarks()) + alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) #interpolate landmarks diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 5750a71..a7a6fa1 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -7,7 +7,7 @@ from pathlib import Path import numpy as np import cv2 from utils import Path_utils -from utils.DFLPNG import DFLPNG +from utils.DFLJPG import DFLJPG from utils import image_utils from facelib import FaceType import facelib @@ -311,7 +311,7 @@ class ExtractSubprocessor(SubprocessorBase): debug_image = image.copy() for (face_idx, face) in enumerate(faces): - output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.png') + output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.jpg') rect = face[0] image_landmarks = np.array(face[1]) @@ -329,7 +329,7 @@ class ExtractSubprocessor(SubprocessorBase): cv2.imwrite(output_file, face_image) - DFLPNG.embed_data(output_file, face_type = FaceType.toString(self.face_type), + DFLJPG.embed_data(output_file, face_type = FaceType.toString(self.face_type), landmarks = face_image_landmarks.tolist(), yaw_value = facelib.LandmarksProcessor.calc_face_yaw (face_image_landmarks), pitch_value = facelib.LandmarksProcessor.calc_face_pitch (face_image_landmarks), diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index f723704..798e4b7 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -10,6 +10,7 @@ from pathlib import Path from utils import Path_utils from utils import image_utils from utils.DFLPNG import DFLPNG +from utils.DFLJPG import DFLJPG from facelib import LandmarksProcessor from utils.SubprocessorBase import SubprocessorBase import multiprocessing @@ -86,17 +87,24 @@ class BlurEstimatorSubprocessor(SubprocessorBase): #override def onClientProcessData(self, data): - filename_path = Path( data[0] ) - - dflpng = DFLPNG.load( str(filename_path), print_on_no_embedded_data=True ) - if dflpng is not None: - image = cv2.imread( str(filename_path) ) - image = ( image * \ - LandmarksProcessor.get_image_hull_mask (image.shape, dflpng.get_landmarks()) \ - ).astype(np.uint8) - return [ str(filename_path), estimate_sharpness( image ) ] + filepath = Path( data[0] ) + + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) else: - return [ str(filename_path), 0 ] + print ("%s is not a dfl image file" % (filepath.name) ) + dflimg = None + + if dflimg is not None: + image = cv2.imread( str(filepath) ) + image = ( image * \ + LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) \ + ).astype(np.uint8) + return [ str(filepath), estimate_sharpness( image ) ] + else: + return [ str(filepath), 0 ] #override def onClientGetDataName (self, data): @@ -151,15 +159,15 @@ def sort_by_face(input_path): for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): filepath = Path(filepath) - if filepath.suffix != '.png': - print ("%s is not a png file required for sort_by_face" % (filepath.name) ) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) continue - - dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) - if dflpng is None: - continue - - img_list.append( [str(filepath), dflpng.get_landmarks()] ) + + img_list.append( [str(filepath), dflimg.get_landmarks()] ) img_list_len = len(img_list) @@ -187,15 +195,15 @@ def sort_by_face_dissim(input_path): for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): filepath = Path(filepath) - if filepath.suffix != '.png': - print ("%s is not a png file required for sort_by_face_dissim" % (filepath.name) ) - continue + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) + continue - dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) - if dflpng is None: - continue - - img_list.append( [str(filepath), dflpng.get_landmarks(), 0 ] ) + img_list.append( [str(filepath), dflimg.get_landmarks(), 0 ] ) img_list_len = len(img_list) for i in tqdm( range(0, img_list_len-1), desc="Sorting", ascii=True): @@ -220,15 +228,15 @@ def sort_by_face_yaw(input_path): for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): filepath = Path(filepath) - if filepath.suffix != '.png': - print ("%s is not a png file required for sort_by_face_dissim" % (filepath.name) ) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) continue - dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) - if dflpng is None: - continue - - img_list.append( [str(filepath), np.array( dflpng.get_yaw_value() ) ] ) + img_list.append( [str(filepath), np.array( dflimg.get_yaw_value() ) ] ) print ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) @@ -436,15 +444,22 @@ def sort_by_hist_dissim(input_path): print ("Sorting by histogram dissimilarity...") img_list = [] - for filename_path in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): - image = cv2.imread(filename_path) + for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): + filepath = Path(filepath) - dflpng = DFLPNG.load( str(filename_path) ) - if dflpng is not None: - face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflpng.get_landmarks()) - image = (image*face_mask).astype(np.uint8) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) + continue + + image = cv2.imread(str(filepath)) + face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) + image = (image*face_mask).astype(np.uint8) - img_list.append ([filename_path, cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ]) + img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ]) img_list = HistDissimSubprocessor(img_list).process() @@ -512,26 +527,27 @@ class FinalLoaderSubprocessor(SubprocessorBase): filepath = Path(data[0]) try: - if filepath.suffix != '.png': - raise Exception ("%s is not a png file required for sort_final" % (filepath.name) ) - - dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) - if dflpng is None: - raise Exception ("") + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) + raise Exception("") bgr = cv2.imread(str(filepath)) if bgr is None: raise Exception ("Unable to load %s" % (filepath.name) ) gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) - gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr.shape, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8) + gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr.shape, dflimg.get_landmarks() )[:,:,0] ).astype(np.uint8) sharpness = estimate_sharpness(gray_masked) hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) except Exception as e: print (e) return [ 1, [str(filepath)] ] - return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ] + return [ 0, [str(filepath), sharpness, hist, dflimg.get_yaw_value() ] ] #override @@ -689,15 +705,15 @@ def sort_by_origname(input_path): for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True): filepath = Path(filepath) - if filepath.suffix != '.png': - print ("%s is not a png file required for sort_by_origname" % (filepath.name) ) - continue - - dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) - if dflpng is None: + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load( str(filepath), print_on_no_embedded_data=True ) + else: + print ("%s is not a dfl image file" % (filepath.name) ) continue - img_list.append( [str(filepath), dflpng.get_source_filename()] ) + img_list.append( [str(filepath), dflimg.get_source_filename()] ) print ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1)) diff --git a/mainscripts/Util.py b/mainscripts/Util.py new file mode 100644 index 0000000..c2fbea1 --- /dev/null +++ b/mainscripts/Util.py @@ -0,0 +1,46 @@ +import os +import sys +import operator +import numpy as np +import cv2 +from tqdm import tqdm +from shutil import copyfile + +from pathlib import Path +from utils import Path_utils +from utils import image_utils +from utils.DFLPNG import DFLPNG +from utils.DFLJPG import DFLJPG +from facelib import LandmarksProcessor +from utils.SubprocessorBase import SubprocessorBase +import multiprocessing + +def convert_png_to_jpg_file (filepath): + filepath = Path(filepath) + + if filepath.suffix != '.png': + return + dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) + if dflpng is None: + return + + dfl_dict = dflpng.getDFLDictData() + + img = cv2.imread (str(filepath)) + new_filepath = str(filepath.parent / (filepath.stem + '.jpg')) + cv2.imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 85]) + DFLJPG.embed_data( new_filepath, **dfl_dict ) + filepath.unlink() + + +def convert_png_to_jpg_folder (input_path): + if not all(ord(c) < 128 for c in input_path): + print ("Path to directory must contain only non unicode characters.") + return + input_path = Path(input_path) + + print ("Converting PNG to JPG...\r\n") + + for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Converting", ascii=True): + filepath = Path(filepath) + convert_png_to_jpg_file(filepath) diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index b1f16d5..06cfaea 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -111,6 +111,7 @@ keras_contrib = nnlib.keras_contrib GroupNormalization = keras_contrib.layers.GroupNormalization InstanceNormalization = keras_contrib.layers.InstanceNormalization Padam = keras_contrib.optimizers.Padam +PELU = keras_contrib.layers.advanced_activations.PELU """ code_import_dlib_string = \ """ @@ -448,20 +449,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator else: return (1.0 - tf.image.ssim ((y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0 nnlib.DSSIMLoss = DSSIMLoss - - class DSSIMLoss(object): - def __init__(self, is_tanh=False): - self.is_tanh = is_tanh - - def __call__(self,y_true, y_pred): - - if not self.is_tanh: - loss = (1.0 - tf.image.ssim (y_true, y_pred, 1.0)) / 2.0 - else: - loss = (1.0 - tf.image.ssim ( (y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0 - - return loss - nnlib.DSSIMLoss = DSSIMLoss class DSSIMMaskLoss(object): def __init__(self, mask_list, is_tanh=False): diff --git a/samples/SampleLoader.py b/samples/SampleLoader.py index 4841d4c..2b5e165 100644 --- a/samples/SampleLoader.py +++ b/samples/SampleLoader.py @@ -6,6 +6,7 @@ from pathlib import Path from utils import Path_utils from utils.DFLPNG import DFLPNG +from utils.DFLJPG import DFLJPG from .Sample import Sample from .Sample import SampleType @@ -54,21 +55,22 @@ class SampleLoader: sample_list = [] for s in tqdm( samples, desc="Loading", ascii=True ): - s_filename_path = Path(s.filename) - if s_filename_path.suffix != '.png': - print ("%s is not a png file required for training" % (s_filename_path.name) ) - continue - - dflpng = DFLPNG.load ( str(s_filename_path), print_on_no_embedded_data=True ) - if dflpng is None: + if s_filename_path.suffix == '.png': + dflimg = DFLPNG.load ( str(s_filename_path), print_on_no_embedded_data=True ) + if dflimg is None: continue + elif s_filename_path.suffix == '.jpg': + dflimg = DFLJPG.load ( str(s_filename_path), print_on_no_embedded_data=True ) + if dflimg is None: continue + else: + print ("%s is not a dfl image file required for training" % (s_filename_path.name) ) continue sample_list.append( s.copy_and_set(sample_type=SampleType.FACE, - face_type=FaceType.fromString (dflpng.get_face_type()), - shape=dflpng.get_shape(), - landmarks=dflpng.get_landmarks(), - yaw=dflpng.get_yaw_value()) ) + face_type=FaceType.fromString (dflimg.get_face_type()), + shape=dflimg.get_shape(), + landmarks=dflimg.get_landmarks(), + yaw=dflimg.get_yaw_value()) ) return sample_list diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py new file mode 100644 index 0000000..eeb1d49 --- /dev/null +++ b/utils/DFLJPG.py @@ -0,0 +1,229 @@ +import struct +import pickle +import numpy as np +from facelib import FaceType +from utils.struct_utils import * + +class DFLJPG(object): + def __init__(self): + self.data = b"" + self.length = 0 + self.chunks = [] + self.dfl_dict = None + + @staticmethod + def load_raw(filename): + try: + with open(filename, "rb") as f: + data = f.read() + except: + raise FileNotFoundError(data) + + 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))) + + return None + + @staticmethod + def load(filename, print_on_no_embedded_data=False, throw_on_no_embedded_data=False): + inst = DFLJPG.load_raw (filename) + 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: + raise Exception("JPG must be in pixel units.") + inst.shape = (Ydensity, Xdensity, 3) + else: + raise Exception("Unknown jpeg ID: %s" % (id) ) + + if chunk['name'] == 'APP15': + if type(chunk['data']) == bytes: + inst.dfl_dict = pickle.loads(chunk['data']) + + if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()): + inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) + + if inst.dfl_dict == None: + if print_on_no_embedded_data: + print ( "No DFL data found in %s" % (filename) ) + if throw_on_no_embedded_data: + raise ValueError("No DFL data found in %s" % (filename) ) + return None + + return inst + + @staticmethod + def embed_data(filename, face_type=None, + landmarks=None, + yaw_value=None, + pitch_value=None, + source_filename=None, + source_rect=None, + source_landmarks=None + ): + + inst = DFLJPG.load_raw (filename) + inst.setDFLDictData ({ + 'face_type': face_type, + 'landmarks': landmarks, + 'yaw_value': yaw_value, + 'pitch_value': pitch_value, + 'source_filename': source_filename, + 'source_rect': source_rect, + 'source_landmarks': source_landmarks + }) + + try: + with open(filename, "wb") as f: + f.write ( inst.dump() ) + except: + raise Exception( 'cannot save %s' % (filename) ) + + 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_yaw_value(self): return self.dfl_dict['yaw_value'] + def get_pitch_value(self): return self.dfl_dict['pitch_value'] + 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'] ) diff --git a/utils/struct_utils.py b/utils/struct_utils.py new file mode 100644 index 0000000..b12bacc --- /dev/null +++ b/utils/struct_utils.py @@ -0,0 +1,6 @@ +import struct + +def struct_unpack(data, counter, fmt): + fmt_size = struct.calcsize(fmt) + return (counter+fmt_size,) + struct.unpack (fmt, data[counter:counter+fmt_size]) + \ No newline at end of file