107 lines
4.6 KiB
Python
107 lines
4.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from abc import ABCMeta
|
|
|
|
import numpy as np
|
|
|
|
import cv2
|
|
import torch
|
|
from PIL import Image
|
|
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
|
from scepter.modules.annotator.registry import ANNOTATORS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
|
|
|
|
def dilate_mask(mask, dilate_factor=15):
|
|
mask = mask.astype(np.uint8)
|
|
mask = cv2.dilate(mask,
|
|
np.ones((dilate_factor, dilate_factor), np.uint8),
|
|
iterations=1)
|
|
return mask
|
|
|
|
|
|
@ANNOTATORS.register_class()
|
|
class LamaAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
|
para_dict = {}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
from modelscope.pipelines.builder import PIPELINES
|
|
from modelscope.pipelines.cv import ImageInpaintingPipeline
|
|
from modelscope.pipelines import pipeline
|
|
from modelscope.utils.constant import Tasks
|
|
from modelscope.metainfo import Pipelines
|
|
from modelscope.models.cv.image_inpainting.refinement import refine_predict
|
|
from torch.utils.data._utils.collate import default_collate
|
|
|
|
@PIPELINES.register_module(Tasks.image_inpainting,
|
|
module_name=Pipelines.image_inpainting +
|
|
'-v2')
|
|
class ImageInpaintingPipelineV2(ImageInpaintingPipeline):
|
|
def perform_inference(self, data):
|
|
px_budget = 9000000
|
|
batch = default_collate([data])
|
|
if self.refine:
|
|
assert 'unpad_to_size' in batch, 'Unpadded size is required for the refinement'
|
|
assert 'cuda' in str(
|
|
self.device), 'GPU is required for refinement'
|
|
gpu_ids = str(self.device).split(':')[-1]
|
|
cur_res = refine_predict(batch,
|
|
self.infer_model,
|
|
gpu_ids=gpu_ids,
|
|
modulo=self.pad_out_to_modulo,
|
|
n_iters=15,
|
|
lr=0.002,
|
|
min_side=512,
|
|
max_scales=3,
|
|
px_budget=px_budget)
|
|
cur_res = cur_res[0].permute(1, 2,
|
|
0).detach().cpu().numpy()
|
|
else:
|
|
with torch.no_grad():
|
|
batch = self.move_to_device(batch, self.device)
|
|
batch['mask'] = (batch['mask'] > 0) * 1
|
|
batch = self.infer_model(batch)
|
|
cur_res = batch['inpainted'][0].permute(
|
|
1, 2, 0).detach().cpu().numpy()
|
|
unpad_to_size = batch.get('unpad_to_size', None)
|
|
if unpad_to_size is not None:
|
|
orig_height, orig_width = unpad_to_size
|
|
cur_res = cur_res[:orig_height, :orig_width]
|
|
|
|
cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
|
|
cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
|
|
return cur_res
|
|
|
|
lama_model_dir = FS.get_dir_to_local_dir(cfg.PRETRAINED_MODEL)
|
|
self.lama_model = pipeline(Tasks.image_inpainting,
|
|
model=lama_model_dir,
|
|
pipeline_name=Pipelines.image_inpainting +
|
|
'-v2',
|
|
refine=True,
|
|
device='cuda:{}'.format(we.device_id))
|
|
|
|
def forward(self, image, mask):
|
|
mask = dilate_mask(mask, dilate_factor=19)
|
|
input_mask = Image.fromarray(mask)
|
|
mask_expanded = np.tile(np.expand_dims(mask, axis=-1), (1, 1, 3))
|
|
input_image_np = np.array(image)
|
|
input_image_np[mask_expanded == 255] = 0
|
|
input_image = Image.fromarray(input_image_np)
|
|
input = {
|
|
'img': input_image,
|
|
'mask': input_mask,
|
|
}
|
|
result = self.lama_model(input)
|
|
output_img = result['output_img']
|
|
return output_img[..., ::-1]
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('ANNOTATORS',
|
|
__class__.__name__,
|
|
LamaAnnotator.para_dict,
|
|
set_name=True)
|