96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
import os
|
|
import io
|
|
import json
|
|
import numpy as np
|
|
import cv2
|
|
import re
|
|
import modules.scripts as scripts
|
|
import modules.shared as shared
|
|
from modules import images
|
|
|
|
from scripts.td_abg import get_foreground
|
|
from scripts.convertor import pil2cv
|
|
from scripts.batch_dir import save_image_dir, modify_basename
|
|
|
|
try:
|
|
from modules.paths_internal import extensions_dir
|
|
except Exception:
|
|
from modules.extensions import extensions_dir
|
|
|
|
from collections import OrderedDict
|
|
from PIL import Image
|
|
|
|
model_cache = OrderedDict()
|
|
sam_model_dir = os.path.join(
|
|
extensions_dir, "PBRemTools/models/")
|
|
model_list = [f for f in os.listdir(sam_model_dir) if os.path.isfile(
|
|
os.path.join(sam_model_dir, f)) and f.split('.')[-1] != 'txt']
|
|
|
|
def processing(single_image, batch_image, input_dir, output_dir, output_mask_dir, show_result, input_tab_state, *rem_args):
|
|
# 0: single
|
|
if (input_tab_state == 0):
|
|
processed = process_image(single_image, *rem_args)
|
|
return processed
|
|
# 1: batch
|
|
elif (input_tab_state == 1):
|
|
processed = []
|
|
for i in batch_image:
|
|
image = Image.open(i)
|
|
base, mask = process_image(image, *rem_args)
|
|
processed.append(base)
|
|
processed.append(mask)
|
|
return processed
|
|
# 2: batch dir (or other)
|
|
else:
|
|
processed = []
|
|
files = shared.listfiles(input_dir)
|
|
for f in files:
|
|
try:
|
|
image = Image.open(f)
|
|
except Exception:
|
|
continue
|
|
base, mask = process_image(image, *rem_args)
|
|
processed.append(base)
|
|
processed.append(mask)
|
|
if output_dir != "":
|
|
basename = os.path.splitext(os.path.basename(f))[0]
|
|
ext = os.path.splitext(f)[1][1:]
|
|
save_image_dir(
|
|
Image.fromarray(base),
|
|
path=output_dir,
|
|
basename=basename,
|
|
extension="png",
|
|
)
|
|
if output_mask_dir != "":
|
|
basename = modify_basename(basename)
|
|
save_image_dir(
|
|
Image.fromarray(mask),
|
|
path=output_mask_dir,
|
|
basename=basename,
|
|
extension="png",
|
|
)
|
|
if (show_result):
|
|
return processed
|
|
else:
|
|
return None
|
|
|
|
def process_image(target_image, *rem_args):
|
|
image = pil2cv(target_image)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
mask, image = get_foreground(image, *rem_args)
|
|
return image, mask
|
|
|
|
class Script(scripts.Script):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def title(self):
|
|
return "PBRemTools"
|
|
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible
|
|
|
|
def ui(self, is_img2img):
|
|
return ()
|
|
|