|
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] |
|
|
