273 lines
12 KiB
Python
273 lines
12 KiB
Python
import os
|
|
import json
|
|
import base64
|
|
import io
|
|
import requests
|
|
from PIL import Image
|
|
import numpy as np
|
|
import argparse
|
|
import random
|
|
from tqdm import tqdm
|
|
import os.path as osp
|
|
from omegaconf import OmegaConf
|
|
from utils.io_utils import find_all_imgs, submit_request, img2b64, save_encoded_image
|
|
from random import randint
|
|
from requests.auth import HTTPBasicAuth
|
|
import cv2
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
import math
|
|
|
|
|
|
INPAINTING_FILL_METHODS = ['fill', 'original', 'latent_noise', 'latent_nothing']
|
|
|
|
|
|
def run_sdinpaint(img: Image.Image, mask: Image.Image, data: dict, prompt: str = '', nprompt: str = '', url='', auth=None) -> str:
|
|
if isinstance(img, Image.Image):
|
|
img_b64 = img2b64(img)
|
|
else:
|
|
assert isinstance(img, str)
|
|
img_b64 = img
|
|
mask_b64 = img2b64(mask)
|
|
data['init_images'] = [img_b64]
|
|
data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64
|
|
data['mask'] = mask_b64
|
|
data['negative_prompt'] = nprompt
|
|
data['prompt'] = prompt
|
|
|
|
response = submit_request(url, json.dumps(data), auth=auth)
|
|
img_b64 = response.json()['images'][0]
|
|
return img_b64
|
|
|
|
def long_side_to(H, W, long_side):
|
|
asp = H / W
|
|
if asp > 1:
|
|
H = long_side
|
|
H = int(round(H / 32)) * 32
|
|
W = int(round(H / asp / 32)) * 32
|
|
else:
|
|
W = long_side
|
|
W = int(round(W / 32)) * 32
|
|
H = int(round(W * asp / 32)) * 32
|
|
return H, W
|
|
|
|
if __name__ == '__main__':
|
|
|
|
parser = argparse.ArgumentParser(description='Inpaint instances of people using stable '
|
|
'diffusion.')
|
|
# workspace\forpaper\eval\kenburns\original\1.png
|
|
parser.add_argument('--img_path', type=str, help='Path to input image. Can be directory.')
|
|
# parser.add_argument('--img_path', type=str, default=r'workspace/style/original/', help='Path to input image.')
|
|
parser.add_argument('--onebyone', type=bool, default=True, help='repainting person one by one')
|
|
parser.add_argument('-n', '--negative_prompt', type=str, default='',
|
|
help='Stable diffusion negative prompt.')
|
|
parser.add_argument('-W', '--width', type=int, default=768, help='Width of output image.')
|
|
parser.add_argument('-H', '--height', type=int, default=768, help='Height of output image.')
|
|
parser.add_argument('-s', '--steps', type=int, default=24, help='Number of diffusion steps.')
|
|
parser.add_argument('-c', '--cfg_scale', type=int, default=9, help='Classifier free guidance '
|
|
'scale, i.e. how strongly the image should conform to prompt.')
|
|
parser.add_argument('-S', '--sample_name', type=str, default='Euler a', help='Name of sampler '
|
|
'to use.')
|
|
parser.add_argument('-d', '--denoising_strength', type=float, default=0.75, help='How much to '
|
|
'disregard original image.')
|
|
parser.add_argument('-f', '--fill', type=str, default=INPAINTING_FILL_METHODS[1],
|
|
help='The fill method to use for inpainting.')
|
|
parser.add_argument('-b', '--mask_blur', type=int, default=4, help='Blur radius of Gaussian '
|
|
'filter to apply to mask.')
|
|
parser.add_argument('-r', '--resolution', type=int, default=640, help='inpainting resolution')
|
|
parser.add_argument('--save_dir', type=str, default='repaint_output')
|
|
parser.add_argument('--url', type=str, default='http://127.0.0.1:7860/sdapi/v1/txt2img', help='img2img url')
|
|
parser.add_argument('--cfg', type=str, default='configs/3d_pixar.yaml', help='repaint config path')
|
|
parser.add_argument('--bg_nprompt', type=str, default='', help='background negative prompt')
|
|
parser.add_argument('--inpaint_full_res', type=int, default=1)
|
|
parser.add_argument('--inpaint_full_res_padding', type=int, default=32)
|
|
parser.add_argument('--detector_ckpt', type=str, default='models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt')
|
|
parser.add_argument('--save_intermediate', type=bool, default=False)
|
|
parser.add_argument('--to-grey', type=bool, default=False)
|
|
parser.add_argument('--infer-tagger', type=bool, default=True)
|
|
parser.add_argument('--style-prompt', default='')
|
|
parser.add_argument('--global-nprompt', default='')
|
|
parser.add_argument('--apply-bg-tagger', default=False)
|
|
parser.add_argument('--apply-fg-tagger', default=True)
|
|
|
|
args = parser.parse_args()
|
|
|
|
args = OmegaConf.create(vars(args))
|
|
args.merge_with(OmegaConf.load(args.cfg))
|
|
|
|
data = {
|
|
**OmegaConf.to_container(args.sd_params),
|
|
# "init_images": [img_b64]
|
|
}
|
|
|
|
auth = None
|
|
if 'username' in args:
|
|
username = args.pop('username')
|
|
password = args.pop('password')
|
|
auth = HTTPBasicAuth(username, password)
|
|
|
|
img_path = args.img_path
|
|
if osp.isfile(img_path):
|
|
imglist = [img_path]
|
|
else:
|
|
imglist = find_all_imgs(img_path, abs_path=True)
|
|
imglist = imglist[::-1]
|
|
|
|
detector = None
|
|
for ii, img_path in enumerate(imglist):
|
|
print(f'repainting {img_path} ... {ii+1}/{len(imglist)}')
|
|
imname = osp.basename(img_path).replace(Path(img_path).suffix, '')
|
|
cimg = Image.open(img_path).convert('RGB')
|
|
W, H = cimg.width, cimg.height
|
|
H, W = long_side_to(H, W, args.long_side)
|
|
data['width'], data['height'] = W, H
|
|
img_resized = cimg.resize((W, H), resample=Image.Resampling.LANCZOS)
|
|
|
|
|
|
if not osp.exists(args.save_dir):
|
|
os.makedirs(args.save_dir)
|
|
|
|
if args.onebyone:
|
|
|
|
repaint_args = {
|
|
'mask_blur': args.mask_blur,
|
|
'inpainting_fill': INPAINTING_FILL_METHODS.index(args.fill),
|
|
'inpaint_full_res': args.inpaint_full_res,
|
|
'inpaint_full_res_padding': args.inpaint_full_res_padding,
|
|
'denoising_strength': args.denoising_strength
|
|
}
|
|
data_inpaint = deepcopy(data)
|
|
data_inpaint.update(repaint_args)
|
|
|
|
from utils.io_utils import json2dict, dict2json, find_all_imgs
|
|
|
|
cache_masks_dir = args.cache_masks_dir
|
|
|
|
mask_fg = None
|
|
masks = []
|
|
|
|
if not osp.exists(cache_masks_dir):
|
|
os.makedirs(cache_masks_dir)
|
|
|
|
|
|
promptp = osp.join(cache_masks_dir, f'{imname}_prompts.json')
|
|
bg_prompt = ''
|
|
fg_prompts = []
|
|
if not osp.exists(promptp):
|
|
if detector is None:
|
|
from animeinsseg import AnimeInsSeg
|
|
import numpy as np
|
|
from animeinsseg.inpainting import patch_match
|
|
detector = AnimeInsSeg(args.detector_ckpt, device='cuda')
|
|
detector.init_tagger()
|
|
instances = detector.infer(img_path, output_type='numpy', infer_tags=True)
|
|
if not instances.is_empty:
|
|
prompts_dict = {}
|
|
for ii, mask in enumerate(instances.masks):
|
|
mask = cv2.resize(mask.astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA)
|
|
mask = Image.fromarray(mask)
|
|
savename = imname + '_' + str(ii).zfill(3) + '.png'
|
|
mask.save(osp.join(cache_masks_dir, savename))
|
|
masks.append(mask)
|
|
tags = instances.tags[ii].split(' ')
|
|
ctags = instances.character_tags[ii]
|
|
for ctag in ctags:
|
|
if ctag in tags:
|
|
tags.remove(ctag)
|
|
prompt = ','.join(tags).replace('_', ' ')
|
|
prompts_dict[savename] = prompt
|
|
fg_prompts.append(prompt)
|
|
|
|
mask_fg = cv2.resize(instances.compose_masks().astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA)
|
|
bg = patch_match.inpaint(np.array(img_resized), mask_fg, patch_size=3)
|
|
savep = osp.join(cache_masks_dir, f'{imname}_bg_repaint.png')
|
|
Image.fromarray(bg).save(savep)
|
|
mask_fg = Image.fromarray(mask_fg)
|
|
mask_fg.save(osp.join(cache_masks_dir, f'{imname}_mask_fg.png'))
|
|
bg_tags, character_tags = detector.tagger.label_cv2_bgr(cv2.cvtColor(bg, cv2.COLOR_BGR2RGB))
|
|
for ii, t in enumerate(bg_tags):
|
|
bg_tags[ii] = t.replace('_', ' ')
|
|
bg_prompt = ','.join(bg_tags)
|
|
prompts_dict[f'{imname}_bg_repaint.png'] = bg_prompt
|
|
dict2json(prompts_dict, promptp)
|
|
else:
|
|
maskp_list = find_all_imgs(cache_masks_dir, abs_path=False)
|
|
prompts_dict = json2dict(promptp)
|
|
for maskn in prompts_dict.keys():
|
|
maskp = osp.join(cache_masks_dir, maskn)
|
|
mask = Image.open(maskp)
|
|
if maskn.endswith('bg_repaint.png'):
|
|
bg_prompt = prompts_dict[maskn]
|
|
bg = mask
|
|
else:
|
|
mask = mask.convert('L')
|
|
fg_prompts.append(prompts_dict[maskn])
|
|
masks.append(mask)
|
|
mask_fg = osp.join(cache_masks_dir, f'{imname}_mask_fg.png')
|
|
mask_fg = Image.open(mask_fg).convert('L')
|
|
|
|
if len(masks) == 0:
|
|
print('no fg is found')
|
|
continue
|
|
|
|
for ii in tqdm(range(args.niter)):
|
|
if args.random_seed:
|
|
data['seed'] = randint(0, 65536)
|
|
else:
|
|
data['seed'] += ii
|
|
seed = data['seed']
|
|
|
|
if args.onebyone:
|
|
data_inpaint['seed'] = seed
|
|
if ii == 0:
|
|
img_b64 = img2b64(bg)
|
|
img_repainted = img_resized
|
|
else:
|
|
img_b64 = output_img_b64
|
|
|
|
if ii == 0:
|
|
nprompt = args.bg_nprompt
|
|
prompt = args.style_prompt + ','
|
|
if args.apply_bg_tagger:
|
|
prompt += bg_prompt + ','
|
|
prompt = prompt.strip(',')
|
|
|
|
data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64
|
|
data['init_images'] = [img_b64]
|
|
data['negative_prompt'] = nprompt
|
|
data['prompt'] = prompt
|
|
response = submit_request(args.url, json.dumps(data), auth=auth)
|
|
output_img_b64 = response.json()['images'][0]
|
|
bg_repainted = Image.open(io.BytesIO(base64.b64decode(output_img_b64)))
|
|
# bg_repainted.show()
|
|
img_repainted = Image.composite(img_repainted, bg_repainted, mask_fg)
|
|
# img_repainted.show()
|
|
|
|
for jj, (fg_prompt, mask) in enumerate(zip(fg_prompts, masks)):
|
|
nprompt = args.global_nprompt
|
|
prompt = args.style_prompt + ','
|
|
if args.apply_fg_tagger:
|
|
prompt += fg_prompt + ','
|
|
print(prompt)
|
|
prompt = prompt.strip(',')
|
|
|
|
output_img_b64 = run_sdinpaint(img_repainted, mask, data_inpaint, prompt=prompt, nprompt=nprompt, url=args.url, auth=auth)
|
|
img_repainted = Image.open(io.BytesIO(base64.b64decode(output_img_b64)))
|
|
# img_repainted.show()
|
|
|
|
save_encoded_image(output_img_b64, osp.join(args.save_dir, f'{imname}_onebyone_niter{ii}_output_{seed}.png'))
|
|
else:
|
|
img_b64 = img2b64(cimg)
|
|
data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64
|
|
data['init_images'] = [img_b64]
|
|
|
|
prompt = args.style_prompt + ','
|
|
prompt = prompt.strip(',')
|
|
data['prompt'] = prompt
|
|
data['negative_prompt'] = args.global_nprompt
|
|
|
|
response = submit_request(args.url, json.dumps(data), auth=auth)
|
|
output_img_b64 = response.json()['images'][0]
|
|
imgsavep = osp.join(args.save_dir, f'{imname}_niter{ii}_output_{seed}.png')
|
|
save_encoded_image(output_img_b64, imgsavep)
|
|
cimg = Image.open(io.BytesIO(base64.b64decode(output_img_b64))) |