DeepFaceLab/samples/SampleLoader.py

import traceback
from enum import IntEnum
import cv2
import numpy as np
from tqdm import tqdm
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

from facelib import FaceType
from facelib import LandmarksProcessor

class SampleLoader:
    cache = dict()
    
    @staticmethod
    def load(sample_type, samples_path, target_samples_path=None):
        cache = SampleLoader.cache
        
        if str(samples_path) not in cache.keys():
            cache[str(samples_path)] = [None]*SampleType.QTY
            
        datas = cache[str(samples_path)]

        if            sample_type == SampleType.IMAGE:
            if  datas[sample_type] is None:  
                datas[sample_type] = [ Sample(filename=filename) for filename in tqdm( Path_utils.get_image_paths(samples_path), desc="Loading", ascii=True ) ]

        elif          sample_type == SampleType.FACE:
            if  datas[sample_type] is None:  
                datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
        
        elif          sample_type == SampleType.FACE_YAW_SORTED:
            if  datas[sample_type] is None:
                datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
        
        elif          sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:            
            if  datas[sample_type] is None:
                if target_samples_path is None:
                    raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET')
                datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) )
        elif          sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:            
            if  datas[sample_type] is None:
                datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) )
        
        
        return datas[sample_type]
        
    @staticmethod
    def upgradeToFaceSamples ( samples ):
        sample_list = []
        
        for s in tqdm( samples, desc="Loading", ascii=True ):
            s_filename_path = Path(s.filename)
            try:
                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 (dflimg.get_face_type()),
                                                   shape=dflimg.get_shape(), 
                                                   landmarks=dflimg.get_landmarks(),
                                                   yaw=dflimg.get_yaw_value()) )
            except:
                print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
                
        return sample_list 
     
    @staticmethod
    def upgradeToFaceCloseToSelfSamples (samples):
        yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples)
        yaw_samples_len = len(yaw_samples)
        
        sample_list = []
        for i in tqdm( range(yaw_samples_len), desc="Sorting", ascii=True ):
            if yaw_samples[i] is not None:
                for s in yaw_samples[i]:
                    s_t = []
                    
                    for n in range(2000):                    
                        yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 )
                        if yaw_samples[yaw_idx] is None:
                            continue
                            
                        yaw_idx_samples_len = len(yaw_samples[yaw_idx])
                        
                        yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ]
                        if s.filename == yaw_idx_sample.filename:
                            continue
                            
                        s_t.append ( yaw_idx_sample )
                        if len(s_t) >= 50:
                            break
                            
                    if len(s_t) == 0:
                        s_t = [s]
                        
                    sample_list.append( s.copy_and_set(close_target_list = s_t) )

        return sample_list
        
    @staticmethod
    def upgradeToFaceYawSortedSamples( samples ):

        lowest_yaw, highest_yaw = -256, +256      
        gradations = 64
        diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations

        yaws_sample_list = [None]*gradations
        
        for i in tqdm( range(0, gradations), desc="Sorting", ascii=True ):
            yaw = lowest_yaw + i*diff_rot_per_grad
            next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad

            yaw_samples = []
            for s in samples:                
                s_yaw = s.yaw
                if (i == 0            and s_yaw < next_yaw) or \
                   (i  < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
                   (i == gradations-1 and s_yaw >= yaw):
                    yaw_samples.append ( s.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED) )
                    
            if len(yaw_samples) > 0:
                yaws_sample_list[i] = yaw_samples
        
        return yaws_sample_list
        
    @staticmethod
    def upgradeToFaceYawSortedAsTargetSamples (s, t):
        l = len(s)
        if l != len(t):
            raise Exception('upgradeToFaceYawSortedAsTargetSamples() s_len != t_len')
        b = l // 2
        
        s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0  for x in s] ) == 1 )[:,0]
        t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0  for x in t] ) == 1 )[:,0]
        
        new_s = [None]*l    
        
        for t_idx in t_idxs:
            search_idxs = []        
            for i in range(0,l):
                search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i]

            for search_idx in search_idxs:            
                if search_idx in s_idxs:
                    mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) )
                    new_s[t_idx] = [ sample.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED_AS_TARGET,
                                                         mirror=True, 
                                                         yaw=-sample.yaw, 
                                                         landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
                                          for sample in s[search_idx] 
                                        ] if mirrored else s[search_idx]                
                    break
                 
        return new_s