889 lines
37 KiB
Python
889 lines
37 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import numbers
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.transforms as TT
|
|
from PIL import Image
|
|
from scepter.modules.annotator.registry import ANNOTATORS
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.studio.preprocess.processors.base_processor import \
|
|
BaseImageProcessor
|
|
|
|
__all__ = [
|
|
'CenterCrop', 'ChangeSample', 'MaskEditSample', 'SwapSample',
|
|
'CannyExtractor', 'ColorExtractor', 'InfoDrawContourExtractor',
|
|
'DegradationExtractor', 'MidasExtractor', 'DoodleExtractor',
|
|
'GrayExtractor', 'InpaintingExtractor', 'OpenposeExtractor',
|
|
'OutpaintingExtractor', 'InfoDrawContourExtractor', 'ESAMExtractor',
|
|
'InvertExtractor', 'DefaultMaskSample', 'MaskSwapEditSample',
|
|
'OutpaintingResize', 'SwapMaskSwapEditSample', 'SourceMaskSample',
|
|
'InpaintingSourceExtractor', 'LamaExtractor'
|
|
]
|
|
|
|
|
|
class CenterCrop(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def process(self, image, height_ratio, width_ratio):
|
|
if isinstance(image, dict):
|
|
image = image['background']
|
|
w, h = image.size
|
|
output_h_align_height, output_h_align_width = h, int(h / height_ratio *
|
|
width_ratio)
|
|
if output_h_align_height * output_h_align_width <= w * h:
|
|
output_height, output_width = output_h_align_height, output_h_align_width
|
|
else:
|
|
output_height, output_width = int(w / width_ratio *
|
|
height_ratio), w
|
|
image = TT.Resize(max(output_height, output_width))(image)
|
|
image = TT.CenterCrop((output_height, output_width))(image)
|
|
return image
|
|
|
|
def __call__(self, **kwargs):
|
|
use_preview = kwargs.get('use_preview', True)
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
target_image = kwargs.get('preview_target_image',
|
|
None) if use_preview else target_image
|
|
src_mask = kwargs.get('preview_src_mask',
|
|
None) if use_preview else src_mask
|
|
src_image = kwargs.get('preview_src_image',
|
|
None) if use_preview else src_image
|
|
height_ratio = kwargs.get('height_ratio', 1)
|
|
width_ratio = kwargs.get('width_ratio', 1)
|
|
ret_data = {}
|
|
if target_image is not None:
|
|
target_image = self.process(target_image, height_ratio,
|
|
width_ratio)
|
|
ret_data['target_image'] = target_image
|
|
if src_image is not None:
|
|
src_image = self.process(src_image, height_ratio, width_ratio)
|
|
ret_data['src_image'] = src_image
|
|
if src_mask is not None:
|
|
src_mask = self.process(src_mask, height_ratio, width_ratio)
|
|
ret_data['src_mask'] = src_mask
|
|
return ret_data
|
|
|
|
|
|
class ChangeSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption', None)
|
|
ret_data = {}
|
|
if preview_target_image is not None:
|
|
ret_data['target_image'] = preview_target_image['composite']
|
|
if preview_src_image is not None:
|
|
ret_data['src_image'] = preview_src_image['background']
|
|
if preview_src_mask is not None:
|
|
ret_data['src_mask'] = preview_src_mask['layers'][0]
|
|
if preview_caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class SourceMaskSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
# target_image = kwargs.get('target_image', None)
|
|
# src_mask = kwargs.get('src_mask', None)
|
|
# src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption', None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['image']
|
|
prev_src_mask = preview_src_image['mask']
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['image']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = None
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
else:
|
|
prev_src_mask = prev_src_mask.convert('L')
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if prev_target_image is not None:
|
|
ret_data['target_image'] = prev_target_image
|
|
else:
|
|
ret_data['target_image'] = Image.new('RGB', (w, h), (0, 0, 0))
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
class MaskEditSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
# target_image = kwargs.get('target_image', None)
|
|
# src_mask = kwargs.get('src_mask', None)
|
|
# src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption',
|
|
None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['background']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = None
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if prev_target_image is not None:
|
|
prev_target_image = prev_target_image.resize((w, h))
|
|
ret_data['target_image'] = Image.composite(
|
|
prev_target_image, prev_src_image, prev_src_mask)
|
|
else:
|
|
ret_data['target_image'] = Image.new('RGB', (w, h), (0, 0, 0))
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class SwapMaskSwapEditSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption', None)
|
|
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['background']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
# process src image and mask according to mask
|
|
ret_data['target_image'] = prev_target_image
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
prev_src_mask_alter = 255 - np.array(prev_src_mask)
|
|
prev_src_mask_alter = Image.fromarray(
|
|
prev_src_mask_alter.astype(np.uint8))
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if prev_target_image is not None:
|
|
prev_target_image = prev_target_image.resize((w, h))
|
|
ret_data['src_image'] = Image.composite(
|
|
prev_target_image, prev_src_image, prev_src_mask_alter)
|
|
else:
|
|
ret_data['src_image'] = Image.new('RGB', (w, h), (0, 0, 0))
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class MaskSwapEditSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption', None)
|
|
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['background']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
# process src image and mask according to mask
|
|
ret_data['target_image'] = prev_src_image
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if prev_target_image is not None:
|
|
prev_target_image = prev_target_image.resize((w, h))
|
|
ret_data['src_image'] = Image.composite(
|
|
prev_target_image, prev_src_image, prev_src_mask)
|
|
else:
|
|
ret_data['src_image'] = Image.new('RGB', (w, h), (0, 0, 0))
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class SwapSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
# replace file with preview data
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
ret_data = {}
|
|
if target_image is not None and src_image is not None:
|
|
ret_data['target_image'] = src_image
|
|
ret_data['src_mask'] = src_mask
|
|
ret_data['src_image'] = target_image
|
|
return ret_data
|
|
|
|
|
|
class DefaultMaskSample(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, **kwargs):
|
|
use_preview = kwargs.get('use_preview', True)
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
preview_target_image = kwargs.get(
|
|
'preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask',
|
|
None) if use_preview else src_mask
|
|
preview_src_image = kwargs.get('preview_src_image',
|
|
None) if use_preview else src_image
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
ret_data = {}
|
|
default_src_mask = np.zeros_like(np.array(prev_src_mask))
|
|
ret_data['src_image'] = prev_src_image
|
|
ret_data['src_mask'] = Image.fromarray(
|
|
default_src_mask.astype(np.uint8))
|
|
ret_data['target_image'] = preview_target_image
|
|
return ret_data
|
|
|
|
|
|
class PaddingCrop(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def get_caption(self, image, **kwargs):
|
|
return image
|
|
|
|
|
|
class BaseExtractor(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
self.model = cfg.get("MODEL", None)
|
|
self.model_info = {'device': 'offline', 'model': None}
|
|
|
|
def unload_model(self):
|
|
if self.use_device.lower() == 'gpu':
|
|
flag, msg = self.unload_model_gpu()
|
|
else:
|
|
flag, msg = self.unload_model_cpu()
|
|
return flag, msg
|
|
|
|
def unload_model_cpu(self):
|
|
super().unload_model()
|
|
self.model_info['device'] = 'offline'
|
|
del self.model_info['model']
|
|
self.model_info['model'] = None
|
|
return True, ''
|
|
|
|
def unload_model_gpu(self):
|
|
allow_load, msg = self.unload_model_cpu()
|
|
if self.delete_instance:
|
|
self.model_info['device'] = 'offline'
|
|
if self.model_info['model'] is not None:
|
|
self.model_info['model'] = self.model_info['model'].to('cpu')
|
|
del self.model_info['model']
|
|
self.model_info['model'] = None
|
|
elif (isinstance(self.model_info['device'], numbers.Number)
|
|
or str(self.model_info['device']).startswith('cuda')):
|
|
self.model_info['device'] = 'cpu'
|
|
self.model_info['model'] = self.model_info['model'].to('cpu')
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
return True, ''
|
|
|
|
def load_model(self):
|
|
allow_load, msg = super().load_model()
|
|
if not allow_load:
|
|
return allow_load, msg
|
|
if self.use_device.lower() == 'gpu':
|
|
flag, msg = self.load_model_gpu()
|
|
else:
|
|
flag, msg = self.load_model_cpu()
|
|
return flag, msg
|
|
|
|
def load_model_cpu(self):
|
|
if self.model_info['device'] == 'offline':
|
|
self.model_info['device'] = 'cpu'
|
|
self.model_info['model'] = ANNOTATORS.build(self.model)
|
|
elif self.model_info['device'] == 'cpu':
|
|
pass
|
|
return True, ''
|
|
|
|
def load_model_gpu(self):
|
|
if self.model_info['device'] == 'offline':
|
|
self.model_info['device'] = 'cpu'
|
|
self.model_info['model'] = ANNOTATORS.build(self.model).to(
|
|
we.device_id)
|
|
elif self.model_info['device'] == 'cpu':
|
|
self.model_info['device'] = we.device_id
|
|
self.model_info['model'] = self.model_info['model'].to(
|
|
we.device_id)
|
|
return True, ''
|
|
|
|
def model_inference(self, model, image, mask=None, **kwargs):
|
|
image = np.array(image)
|
|
return model(image)
|
|
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get(
|
|
'preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask',
|
|
None) if use_preview else src_mask
|
|
preview_src_image = kwargs.get('preview_src_image',
|
|
None) if use_preview else src_image
|
|
preview_caption = kwargs.get('preview_caption',
|
|
None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background'].convert('RGB')
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = None
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
prev_src_mask = ret_data['src_mask']
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
preview_target_image_np = self.model_inference(
|
|
self.model_info['model'], prev_src_image, prev_src_mask)
|
|
if len(preview_target_image_np.shape) > 3:
|
|
preview_target_image = Image.fromarray(
|
|
preview_target_image_np[0])
|
|
else:
|
|
preview_target_image = Image.fromarray(preview_target_image_np)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if preview_target_image is not None and np.sum(
|
|
np.array(prev_src_mask)) > 0:
|
|
preview_target_image = preview_target_image.resize((w, h))
|
|
ret_data['target_image'] = Image.composite(
|
|
preview_target_image, prev_src_image, prev_src_mask)
|
|
else:
|
|
ret_data['target_image'] = preview_target_image
|
|
elif preview_target_image is not None:
|
|
preview_target_image = Image.fromarray(
|
|
self.model_inference(self.model_info['model'],
|
|
preview_target_image))
|
|
ret_data['target_image'] = preview_target_image
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class CannyExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class ColorExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class InfoDrawContourExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class DegradationExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class MidasExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class DoodleExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class GrayExtractor(BaseExtractor):
|
|
pass
|
|
|
|
class LamaExtractor(BaseExtractor):
|
|
def model_inference(self, model, image, mask=None, **kwargs):
|
|
image = np.array(image)
|
|
mask = np.array(mask)
|
|
return model(image, mask)
|
|
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None)
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption', None)
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['image']
|
|
prev_src_mask = preview_src_image['mask']
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = prev_src_image
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
prev_src_mask = ret_data['src_mask']
|
|
else:
|
|
prev_src_mask = prev_src_mask.convert('L')
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
preview_target_image_np = self.model_inference(
|
|
self.model_info['model'], prev_src_image, prev_src_mask)
|
|
if len(preview_target_image_np.shape) > 3:
|
|
preview_target_image = Image.fromarray(
|
|
preview_target_image_np[0])
|
|
else:
|
|
preview_target_image = Image.fromarray(preview_target_image_np)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
ret_data['target_image'] = preview_target_image
|
|
elif preview_target_image is not None:
|
|
preview_target_image = Image.fromarray(
|
|
self.model_inference(self.model_info['model'],
|
|
preview_target_image))
|
|
ret_data['target_image'] = preview_target_image
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
class InpaintingExtractor(BaseExtractor):
|
|
def model_inference(self, model, image, mask=None, **kwargs):
|
|
image = np.array(image)
|
|
mask = np.array(mask) if mask is not None else None
|
|
return model(image, mask)
|
|
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get(
|
|
'preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption',
|
|
None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if preview_target_image is not None:
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['background']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = None
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
prev_src_mask = ret_data['src_mask']
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
preview_src_image_np = self.model_inference(
|
|
self.model_info['model'], prev_src_image, prev_src_mask)
|
|
if len(preview_src_image_np.shape) > 3:
|
|
preview_src_image = Image.fromarray(preview_src_image_np[0])
|
|
else:
|
|
preview_src_image = Image.fromarray(preview_src_image_np)
|
|
# process target image: target_image * mask + src_image * (1 - mask)
|
|
if preview_src_image is not None and np.sum(
|
|
np.array(prev_src_mask)) > 0:
|
|
preview_src_image = preview_src_image.resize((w, h))
|
|
ret_data['src_image'] = Image.composite(
|
|
preview_src_image, prev_src_image, prev_src_mask)
|
|
else:
|
|
ret_data['src_image'] = preview_src_image
|
|
ret_data['target_image'] = prev_target_image
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class InpaintingSourceExtractor(BaseExtractor):
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
preview_target_image = kwargs.get('preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask', None) if use_preview else src_mask
|
|
preview_src_image = kwargs.get('preview_src_image', None) if use_preview else src_image
|
|
preview_caption = kwargs.get('preview_caption', None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['image']
|
|
prev_src_mask = preview_src_mask['image']
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
|
|
if preview_target_image is not None:
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['image']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
|
|
# process src image and mask according to mask
|
|
ret_data['src_image'] = None
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
src_mask = ret_data['src_mask']
|
|
else:
|
|
src_mask = prev_src_mask.convert('L')
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
ret_data['src_image'] = Image.composite(Image.fromarray(255 - np.array(prev_src_mask)), prev_src_image, src_mask)
|
|
ret_data['target_image'] = prev_target_image
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class OpenposeExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class OutpaintingExtractor(BaseExtractor):
|
|
def model_inference(self,
|
|
model,
|
|
image,
|
|
mask=None,
|
|
use_mask=False,
|
|
**kwargs):
|
|
image = np.array(image)
|
|
return model(image, mask=mask if use_mask else None, return_mask=True)
|
|
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
use_mask = kwargs.get('use_mask', False)
|
|
# preview_target_image = kwargs.get('preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask', None)
|
|
preview_src_image = kwargs.get('preview_src_image', None)
|
|
preview_caption = kwargs.get('preview_caption',
|
|
None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image[
|
|
'background'] if use_preview else src_image
|
|
prev_src_mask = preview_src_image['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_image = preview_src_image if use_preview else src_image
|
|
prev_src_mask = preview_src_mask
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
# process src image and mask according to mask
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
prev_src_mask = ret_data['src_mask']
|
|
else:
|
|
prev_src_mask = prev_src_mask
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
outpaiting_results = self.model_inference(self.model_info['model'],
|
|
prev_src_image,
|
|
prev_src_mask,
|
|
use_mask=use_mask)
|
|
src_image, target_image, src_mask = (
|
|
outpaiting_results['src_image'], outpaiting_results['image'],
|
|
outpaiting_results['mask'])
|
|
if len(target_image.shape) > 3:
|
|
preview_target_image = Image.fromarray(target_image[0])
|
|
else:
|
|
preview_target_image = Image.fromarray(target_image)
|
|
|
|
if len(src_image.shape) > 3:
|
|
preview_src_image = Image.fromarray(src_image[0])
|
|
else:
|
|
preview_src_image = Image.fromarray(src_image)
|
|
|
|
if len(src_mask.shape) > 3:
|
|
preview_src_mask = Image.fromarray(src_mask[0])
|
|
else:
|
|
preview_src_mask = Image.fromarray(src_mask)
|
|
|
|
ret_data['target_image'] = preview_target_image
|
|
ret_data['src_image'] = preview_src_image
|
|
ret_data['src_mask'] = preview_src_mask
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class OutpaintingResize(BaseExtractor):
|
|
def model_inference(self, model, image, target_image, mask=None, **kwargs):
|
|
image = np.array(image)
|
|
return model(image, target_image, mask=mask)
|
|
|
|
def __call__(self, **kwargs):
|
|
target_image = kwargs.get('target_image', None)
|
|
src_mask = kwargs.get('src_mask', None)
|
|
src_image = kwargs.get('src_image', None)
|
|
use_preview = kwargs.get('use_preview', True)
|
|
caption = kwargs.get('caption', None)
|
|
use_mask = kwargs.get('use_mask', False)
|
|
preview_target_image = kwargs.get(
|
|
'preview_target_image', None) if use_preview else target_image
|
|
preview_src_mask = kwargs.get('preview_src_mask',
|
|
None) if use_preview else src_mask
|
|
preview_src_image = kwargs.get('preview_src_image',
|
|
None) if use_preview else src_image
|
|
preview_caption = kwargs.get('preview_caption',
|
|
None) if use_preview else caption
|
|
ret_data = {}
|
|
if preview_src_image is not None:
|
|
if isinstance(preview_src_image, dict):
|
|
prev_src_image = preview_src_image['background']
|
|
else:
|
|
prev_src_image = preview_src_image
|
|
if isinstance(preview_src_mask, dict):
|
|
prev_src_mask = preview_src_mask['layers'][0].split(
|
|
)[-1].convert('L')
|
|
else:
|
|
prev_src_mask = preview_src_mask
|
|
if isinstance(preview_target_image, dict):
|
|
prev_target_image = preview_target_image['background']
|
|
else:
|
|
prev_target_image = preview_target_image
|
|
else:
|
|
prev_src_image = None
|
|
prev_src_mask = None
|
|
prev_target_image = None
|
|
# process src image and mask according to mask
|
|
if prev_src_image is not None:
|
|
w, h = prev_src_image.size
|
|
if prev_src_mask is None:
|
|
ret_data['src_mask'] = Image.new('L', (w, h), 0)
|
|
prev_src_mask = ret_data['src_mask']
|
|
else:
|
|
prev_src_mask = prev_src_mask.convert('L')
|
|
ow, oh = prev_src_mask.size
|
|
if ow == w and oh == h:
|
|
ret_data['src_mask'] = prev_src_mask
|
|
else:
|
|
ret_data['src_mask'] = Image.new('L', (ow, oh), 0)
|
|
outpaiting_results = self.model_inference(self.model_info['model'],
|
|
prev_src_image,
|
|
prev_target_image,
|
|
prev_src_mask)
|
|
src_image = outpaiting_results['src_image']
|
|
if len(src_image.shape) > 3:
|
|
preview_src_image = Image.fromarray(src_image[0])
|
|
else:
|
|
preview_src_image = Image.fromarray(src_image)
|
|
ret_data['src_image'] = preview_src_image
|
|
ret_data['target_image'] = preview_target_image
|
|
ret_data['src_mask'] = preview_src_mask
|
|
if caption is not None:
|
|
ret_data['caption'] = preview_caption
|
|
return ret_data
|
|
|
|
|
|
class InfoDrawAnimeAnnotator(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class ESAMExtractor(BaseExtractor):
|
|
pass
|
|
|
|
|
|
class InvertExtractor(BaseExtractor):
|
|
pass
|