mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-03-12 20:42:45 -07:00
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
import multiprocessing
|
|
import operator
|
|
import traceback
|
|
from pathlib import Path
|
|
import pickle
|
|
import samplelib.PackedFaceset
|
|
from DFLIMG import *
|
|
from facelib import FaceType, LandmarksProcessor
|
|
from interact import interact as io
|
|
from joblib import Subprocessor
|
|
from utils import Path_utils, mp_utils
|
|
|
|
from .Sample import Sample, SampleType
|
|
|
|
|
|
class SampleHost:
|
|
|
|
|
|
|
|
|
|
samples_cache = dict()
|
|
@staticmethod
|
|
def get_person_id_max_count(samples_path):
|
|
samples = None
|
|
try:
|
|
samples = samplelib.PackedFaceset.load(samples_path)
|
|
except:
|
|
io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}")
|
|
|
|
if samples is None:
|
|
raise ValueError("packed faceset not found.")
|
|
persons_name_idxs = {}
|
|
for sample in samples:
|
|
persons_name_idxs[sample.person_name] = 0
|
|
return len(list(persons_name_idxs.keys()))
|
|
|
|
@staticmethod
|
|
def load(sample_type, samples_path):
|
|
samples_cache = SampleHost.samples_cache
|
|
|
|
if str(samples_path) not in samples_cache.keys():
|
|
samples_cache[str(samples_path)] = [None]*SampleType.QTY
|
|
|
|
samples = samples_cache[str(samples_path)]
|
|
|
|
if sample_type == SampleType.IMAGE:
|
|
if samples[sample_type] is None:
|
|
samples[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ]
|
|
|
|
elif sample_type == SampleType.FACE:
|
|
if samples[sample_type] is None:
|
|
try:
|
|
result = samplelib.PackedFaceset.load(samples_path)
|
|
except:
|
|
io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}")
|
|
|
|
if result is not None:
|
|
io.log_info (f"Loaded {len(result)} packed faces from {samples_path}")
|
|
|
|
if result is None:
|
|
result = SampleHost.load_face_samples( Path_utils.get_image_paths(samples_path) )
|
|
samples[sample_type] = result
|
|
|
|
elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
|
|
result = SampleHost.load (SampleType.FACE, samples_path)
|
|
result = SampleHost.upgradeToFaceTemporalSortedSamples(result)
|
|
samples[sample_type] = result
|
|
|
|
return samples[sample_type]
|
|
|
|
@staticmethod
|
|
def load_face_samples ( image_paths):
|
|
sample_list = []
|
|
|
|
for filename in io.progress_bar_generator (image_paths, desc="Loading"):
|
|
dflimg = DFLIMG.load (Path(filename))
|
|
if dflimg is None:
|
|
io.log_err (f"{filename} is not a dfl image file.")
|
|
else:
|
|
sample_list.append( Sample(filename=filename,
|
|
sample_type=SampleType.FACE,
|
|
face_type=FaceType.fromString ( dflimg.get_face_type() ),
|
|
shape=dflimg.get_shape(),
|
|
landmarks=dflimg.get_landmarks(),
|
|
ie_polys=dflimg.get_ie_polys(),
|
|
eyebrows_expand_mod=dflimg.get_eyebrows_expand_mod(),
|
|
source_filename=dflimg.get_source_filename(),
|
|
))
|
|
return sample_list
|
|
|
|
@staticmethod
|
|
def upgradeToFaceTemporalSortedSamples( samples ):
|
|
new_s = [ (s, s.source_filename) for s in samples]
|
|
new_s = sorted(new_s, key=operator.itemgetter(1))
|
|
|
|
return [ s[0] for s in new_s]
|