| import multiprocessing |
| import pickle |
| import time |
| import traceback |
| from enum import IntEnum |
|
|
| import cv2 |
| import numpy as np |
|
|
| from core import imagelib, mplib, pathex |
| from core.imagelib import sd |
| from core.cv2ex import * |
| from core.interact import interact as io |
| from core.joblib import SubprocessGenerator, ThisThreadGenerator |
| from facelib import LandmarksProcessor |
| from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor, SampleType) |
|
|
| class SampleGeneratorFaceAvatarOperator(SampleGeneratorBase): |
| def __init__ (self, root_path, debug=False, batch_size=1, resolution=256, face_type=None, |
| generators_count=4, data_format="NHWC", |
| **kwargs): |
|
|
| super().__init__(debug, batch_size) |
| self.initialized = False |
|
|
|
|
| dataset_path = root_path / 'AvatarOperatorDataset' |
| if not dataset_path.exists(): |
| raise ValueError(f'Unable to find {dataset_path}') |
|
|
| chains_dir_names = pathex.get_all_dir_names(dataset_path) |
|
|
| samples = SampleLoader.load (SampleType.FACE, dataset_path, subdirs=True) |
| sample_idx_by_path = { sample.filename : i for i,sample in enumerate(samples) } |
|
|
| kf_idxs = [] |
|
|
| for chain_dir_name in chains_dir_names: |
| chain_root_path = dataset_path / chain_dir_name |
|
|
| subchain_dir_names = pathex.get_all_dir_names(chain_root_path) |
| try: |
| subchain_dir_names.sort(key=int) |
| except: |
| raise Exception(f'{chain_root_path} must contain only numerical name of directories') |
| chain_samples = [] |
|
|
| for subchain_dir_name in subchain_dir_names: |
| subchain_root = chain_root_path / subchain_dir_name |
| subchain_samples = [ sample_idx_by_path[image_path] for image_path in pathex.get_image_paths(subchain_root) \ |
| if image_path in sample_idx_by_path ] |
|
|
| if len(subchain_samples) < 3: |
| raise Exception(f'subchain {subchain_dir_name} must contain at least 3 faces. If you delete this subchain, then th echain will be corrupted.') |
|
|
| chain_samples += [ subchain_samples ] |
|
|
| chain_samples_len = len(chain_samples) |
| for i in range(chain_samples_len-1): |
| kf_idxs += [ ( chain_samples[i+1][0], chain_samples[i][-1], chain_samples[i][:-1] ) ] |
| |
| for i in range(1,chain_samples_len): |
| kf_idxs += [ ( chain_samples[i-1][-1], chain_samples[i][0], chain_samples[i][1:] ) ] |
|
|
| if self.debug: |
| self.generators_count = 1 |
| else: |
| self.generators_count = max(1, generators_count) |
|
|
| if self.debug: |
| self.generators = [ThisThreadGenerator ( self.batch_func, (samples, kf_idxs, resolution, face_type, data_format) )] |
| else: |
| self.generators = [SubprocessGenerator ( self.batch_func, (samples, kf_idxs, resolution, face_type, data_format), start_now=False ) \ |
| for i in range(self.generators_count) ] |
|
|
| SubprocessGenerator.start_in_parallel( self.generators ) |
|
|
| self.generator_counter = -1 |
|
|
| self.initialized = True |
|
|
| |
| def is_initialized(self): |
| return self.initialized |
|
|
| def __iter__(self): |
| return self |
|
|
| def __next__(self): |
| self.generator_counter += 1 |
| generator = self.generators[self.generator_counter % len(self.generators) ] |
| return next(generator) |
|
|
| def batch_func(self, param ): |
| samples, kf_idxs, resolution, face_type, data_format = param |
| |
| kf_idxs_len = len(kf_idxs) |
|
|
| shuffle_idxs = [] |
| idxs = [*range(len(samples))] |
|
|
| random_flip = True |
| rotation_range=[-10,10] |
| scale_range=[-0.05, 0.05] |
| tx_range=[-0.05, 0.05] |
| ty_range=[-0.05, 0.05] |
|
|
| bs = self.batch_size |
| while True: |
| batches = [ [], [] , [], [], [], [] ] |
|
|
| n_batch = 0 |
| while n_batch < bs: |
| try: |
| if len(shuffle_idxs) == 0: |
| shuffle_idxs = idxs.copy() |
| np.random.shuffle(shuffle_idxs) |
| idx = shuffle_idxs.pop() |
|
|
|
|
| key_idx, key_chain_idx, chain_idxs = kf_idxs[ np.random.randint(kf_idxs_len) ] |
| |
| key_sample = samples[key_idx] |
| key_chain_sample = samples[key_chain_idx] |
| chain_sample = samples[ chain_idxs[np.random.randint(len(chain_idxs)) ] ] |
| |
| |
| |
| |
| |
| |
| sample = samples[idx] |
|
|
| img = sample.load_bgr() |
| |
| key_img = key_sample.load_bgr() |
| key_chain_img = key_chain_sample.load_bgr() |
| chain_img = chain_sample.load_bgr() |
| |
| h,w,c = img.shape |
|
|
| mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks) |
| mask = np.clip(mask, 0, 1) |
| |
| warp_params = imagelib.gen_warp_params(resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range ) |
|
|
| if face_type == sample.face_type: |
| if w != resolution: |
| img = cv2.resize( img, (resolution, resolution), cv2.INTER_CUBIC ) |
| key_img = cv2.resize( key_img, (resolution, resolution), cv2.INTER_CUBIC ) |
| key_chain_img = cv2.resize( key_chain_img, (resolution, resolution), cv2.INTER_CUBIC ) |
| chain_img = cv2.resize( chain_img, (resolution, resolution), cv2.INTER_CUBIC ) |
| |
| mask = cv2.resize( mask, (resolution, resolution), cv2.INTER_CUBIC ) |
| else: |
| mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, face_type) |
| img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) |
| key_img = cv2.warpAffine( key_img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) |
| key_chain_img = cv2.warpAffine( key_chain_img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) |
| chain_img = cv2.warpAffine( chain_img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC ) |
| mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_CUBIC ) |
|
|
| if len(mask.shape) == 2: |
| mask = mask[...,None] |
|
|
| img_warped = imagelib.warp_by_params (warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=True) |
| img_transformed = imagelib.warp_by_params (warp_params, img, can_warp=False, can_transform=True, can_flip=True, border_replicate=True) |
| |
| mask = imagelib.warp_by_params (warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) |
|
|
| key_img = imagelib.warp_by_params (warp_params, key_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) |
| key_chain_img = imagelib.warp_by_params (warp_params, key_chain_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) |
| chain_img = imagelib.warp_by_params (warp_params, chain_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) |
| |
| |
| img_warped = np.clip(img_warped.astype(np.float32), 0, 1) |
| img_transformed = np.clip(img_transformed.astype(np.float32), 0, 1) |
| mask[mask < 0.5] = 0.0 |
| mask[mask >= 0.5] = 1.0 |
| mask = np.clip(mask, 0, 1) |
|
|
| if data_format == "NCHW": |
| img_warped = np.transpose(img_warped, (2,0,1) ) |
| img_transformed = np.transpose(img_transformed, (2,0,1) ) |
| mask = np.transpose(mask, (2,0,1) ) |
| |
| key_img = np.transpose(key_img, (2,0,1) ) |
| key_chain_img = np.transpose(key_chain_img, (2,0,1) ) |
| chain_img = np.transpose(chain_img, (2,0,1) ) |
|
|
| batches[0].append ( img_warped ) |
| batches[1].append ( img_transformed ) |
| batches[2].append ( mask ) |
| batches[3].append ( key_img ) |
| batches[4].append ( key_chain_img ) |
| batches[5].append ( chain_img ) |
|
|
| n_batch += 1 |
| except: |
| io.log_err ( traceback.format_exc() ) |
|
|
| yield [ np.array(batch) for batch in batches] |
|
|
