335 lines
15 KiB
Python
335 lines
15 KiB
Python
import os
|
|
import gc
|
|
import glob
|
|
import copy
|
|
from PIL import Image
|
|
from collections import OrderedDict
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
from segment_anything import SamAutomaticMaskGenerator
|
|
from modules import scripts, shared
|
|
from modules.paths import extensions_dir
|
|
from modules.devices import torch_gc
|
|
|
|
|
|
global_sam: SamAutomaticMaskGenerator = None
|
|
sem_seg_cache = OrderedDict()
|
|
sam_annotator_dir = os.path.join(scripts.basedir(), "annotator")
|
|
original_uniformer_inference_segmentor = None
|
|
|
|
|
|
def pad64(x):
|
|
return int(np.ceil(float(x) / 64.0) * 64 - x)
|
|
|
|
|
|
def safer_memory(x):
|
|
# Fix many MAC/AMD problems
|
|
return np.ascontiguousarray(x.copy()).copy()
|
|
|
|
|
|
def resize_image_with_pad(input_image, resolution):
|
|
from annotator.util import HWC3
|
|
img = HWC3(input_image)
|
|
H_raw, W_raw, _ = img.shape
|
|
k = float(resolution) / float(min(H_raw, W_raw))
|
|
interpolation = cv2.INTER_CUBIC if k > 1 else cv2.INTER_AREA
|
|
H_target = int(np.round(float(H_raw) * k))
|
|
W_target = int(np.round(float(W_raw) * k))
|
|
img = cv2.resize(img, (W_target, H_target), interpolation=interpolation)
|
|
H_pad, W_pad = pad64(H_target), pad64(W_target)
|
|
img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode='edge')
|
|
|
|
def remove_pad(x):
|
|
return safer_memory(x[:H_target, :W_target])
|
|
|
|
return safer_memory(img_padded), remove_pad
|
|
|
|
|
|
def blend_image_and_seg(image, seg, alpha=0.5):
|
|
image_blend = image * (1 - alpha) + np.array(seg) * alpha
|
|
return Image.fromarray(image_blend.astype(np.uint8))
|
|
|
|
|
|
def create_symbolic_link():
|
|
cnet_annotator_dir = os.path.join(extensions_dir, "sd-webui-controlnet/annotator")
|
|
if os.path.isdir(cnet_annotator_dir):
|
|
if not os.path.isdir(sam_annotator_dir):
|
|
os.symlink(cnet_annotator_dir, sam_annotator_dir, target_is_directory=True)
|
|
return True
|
|
return False
|
|
|
|
|
|
def clear_sem_sam_cache():
|
|
sem_seg_cache.clear()
|
|
gc.collect()
|
|
torch_gc()
|
|
|
|
|
|
def sem_sam_garbage_collect():
|
|
if shared.cmd_opts.lowvram:
|
|
for model_key, model in sem_seg_cache:
|
|
if model_key == "uniformer":
|
|
from annotator.uniformer import unload_uniformer_model
|
|
unload_uniformer_model()
|
|
else:
|
|
model.unload_model()
|
|
gc.collect()
|
|
torch_gc()
|
|
|
|
|
|
def strengthen_sem_seg(class_ids, img):
|
|
print("Auto SAM strengthening semantic segmentation")
|
|
import pycocotools.mask as maskUtils
|
|
semantc_mask = copy.deepcopy(class_ids)
|
|
annotations = global_sam.generate(img)
|
|
annotations = sorted(annotations, key=lambda x: x['area'], reverse=True)
|
|
print(f"Auto SAM generated {len(annotations)} masks")
|
|
for ann in annotations:
|
|
valid_mask = torch.tensor(maskUtils.decode(ann['segmentation'])).bool()
|
|
propose_classes_ids = torch.tensor(class_ids[valid_mask])
|
|
num_class_proposals = len(torch.unique(propose_classes_ids))
|
|
if num_class_proposals == 1:
|
|
semantc_mask[valid_mask] = propose_classes_ids[0].numpy()
|
|
continue
|
|
top_1_propose_class_ids = torch.bincount(propose_classes_ids.flatten()).topk(1).indices
|
|
semantc_mask[valid_mask] = top_1_propose_class_ids.numpy()
|
|
print("Auto SAM strengthen process end")
|
|
return semantc_mask
|
|
|
|
|
|
def random_segmentation(img):
|
|
print("Auto SAM generating random segmentation for Edit-Anything")
|
|
img_np = np.array(img.convert("RGB"))
|
|
annotations = global_sam.generate(img_np)
|
|
# annotations = sorted(annotations, key=lambda x: x['area'], reverse=True)
|
|
print(f"Auto SAM generated {len(annotations)} masks")
|
|
H, W, _ = img_np.shape
|
|
color_map = np.zeros((H, W, 3), dtype=np.uint8)
|
|
detected_map_tmp = np.zeros((H, W), dtype=np.uint16)
|
|
for idx, annotation in enumerate(annotations):
|
|
current_seg = annotation['segmentation']
|
|
color_map[current_seg] = np.random.randint(0, 255, (3))
|
|
detected_map_tmp[current_seg] = idx + 1
|
|
detected_map = np.zeros((detected_map_tmp.shape[0], detected_map_tmp.shape[1], 3))
|
|
detected_map[:, :, 0] = detected_map_tmp % 256
|
|
detected_map[:, :, 1] = detected_map_tmp // 256
|
|
from annotator.util import HWC3
|
|
detected_map = HWC3(detected_map.astype(np.uint8))
|
|
print("Auto SAM generation process end")
|
|
return [blend_image_and_seg(img_np, color_map), Image.fromarray(color_map), Image.fromarray(detected_map)], \
|
|
"Random segmentation done. Left above (0) is blended image, right above (1) is random segmentation, left below (2) is Edit-Anything control input."
|
|
|
|
|
|
def image_layer_image(layout_input_image, layout_output_path):
|
|
img_np = np.array(layout_input_image.convert("RGB"))
|
|
annotations = global_sam.generate(img_np)
|
|
print(f"AutoSAM generated {len(annotations)} annotations")
|
|
annotations = sorted(annotations, key=lambda x: x['area'])
|
|
for idx, annotation in enumerate(annotations):
|
|
img_tmp = np.zeros((img_np.shape[0], img_np.shape[1], 3))
|
|
img_tmp[annotation['segmentation']] = img_np[annotation['segmentation']]
|
|
img_np[annotation['segmentation']] = np.array([0, 0, 0])
|
|
img_tmp = Image.fromarray(img_tmp.astype(np.uint8))
|
|
img_tmp.save(os.path.join(layout_output_path, f"{idx}.png"))
|
|
img_np = Image.fromarray(img_np.astype(np.uint8))
|
|
img_np.save(os.path.join(layout_output_path, "leftover.png"))
|
|
|
|
|
|
def image_layer_internal(layout_input_image_or_path, layout_output_path):
|
|
if isinstance(layout_input_image_or_path, str):
|
|
print("Image layer division batch processing")
|
|
all_files = glob.glob(os.path.join(layout_input_image_or_path, "*"))
|
|
for image_index, input_image_file in enumerate(all_files):
|
|
print(f"Processing {image_index}/{len(all_files)} {input_image_file}")
|
|
try:
|
|
input_image = Image.open(input_image_file)
|
|
output_directory = os.path.join(layout_output_path, os.path.splitext(os.path.basename(input_image_file))[0])
|
|
from pathlib import Path
|
|
Path(output_directory).mkdir(exist_ok=True)
|
|
except:
|
|
print(f"File {input_image_file} not image, skipped.")
|
|
continue
|
|
image_layer_image(input_image, output_directory)
|
|
else:
|
|
image_layer_image(layout_input_image_or_path, layout_output_path)
|
|
return "Done"
|
|
|
|
|
|
def inject_uniformer_inference_segmentor(model, img):
|
|
original_result = original_uniformer_inference_segmentor(model, img)
|
|
original_result[0] = strengthen_sem_seg(original_result[0], img)
|
|
return original_result
|
|
|
|
|
|
def inject_uniformer_show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10), opacity=0.5, title='', block=True):
|
|
return result[0]
|
|
|
|
|
|
def inject_oneformer_semantic_run(img, predictor, metadata):
|
|
predictions = predictor(img[:, :, ::-1], "semantic") # Predictor of OneFormer must use BGR image !!!
|
|
original_result = predictions["sem_seg"].argmax(dim=0).cpu().numpy()
|
|
from annotator.oneformer.oneformer.demo.visualizer import Visualizer, ColorMode
|
|
visualizer_map = Visualizer(img, is_img=False, metadata=metadata, instance_mode=ColorMode.IMAGE)
|
|
out_map = visualizer_map.draw_sem_seg(torch.tensor(strengthen_sem_seg(original_result, img), device='cpu'), alpha=1, is_text=False).get_image()
|
|
return out_map
|
|
|
|
|
|
def inject_oneformer_semantic_run_categorical_mask(img, predictor, metadata):
|
|
predictions = predictor(img[:, :, ::-1], "semantic") # Predictor of OneFormer must use BGR image !!!
|
|
original_result = predictions["sem_seg"].argmax(dim=0).cpu().numpy()
|
|
return strengthen_sem_seg(original_result, img)
|
|
|
|
|
|
def inject_oneformer_detector_call(self, img):
|
|
if self.model is None:
|
|
self.load_model()
|
|
self.model.model.to(self.device)
|
|
return inject_oneformer_semantic_run(img, self.model, self.metadata)
|
|
|
|
|
|
def inject_oneformer_detector_call_categorical_mask(self, img):
|
|
if self.model is None:
|
|
self.load_model()
|
|
self.model.model.to(self.device)
|
|
return inject_oneformer_semantic_run_categorical_mask(img, self.model, self.metadata)
|
|
|
|
|
|
def _uniformer(img):
|
|
if "uniformer" not in sem_seg_cache:
|
|
from annotator.uniformer import apply_uniformer
|
|
sem_seg_cache["uniformer"] = apply_uniformer
|
|
result = sem_seg_cache["uniformer"](img)
|
|
return result
|
|
|
|
|
|
def _oneformer(img, dataset="coco"):
|
|
oneformer_key = f"oneformer_{dataset}"
|
|
if oneformer_key not in sem_seg_cache:
|
|
from annotator.oneformer import OneformerDetector
|
|
sem_seg_cache[oneformer_key] = OneformerDetector(OneformerDetector.configs[dataset])
|
|
result = sem_seg_cache[oneformer_key](img)
|
|
return result
|
|
|
|
|
|
def semantic_segmentation(input_image, annotator_name, processor_res,
|
|
use_pixel_perfect, resize_mode, target_W, target_H):
|
|
if input_image is None:
|
|
return [], "No input image."
|
|
if "seg" in annotator_name:
|
|
if not os.path.isdir(os.path.join(scripts.basedir(), "annotator")) and not create_symbolic_link():
|
|
return [], "ControlNet extension not found."
|
|
global original_uniformer_inference_segmentor
|
|
input_image_np = np.array(input_image)
|
|
processor_res = pixel_perfect_lllyasviel(input_image_np, processor_res, use_pixel_perfect, resize_mode, target_W, target_H)
|
|
input_image, remove_pad = resize_image_with_pad(input_image_np, processor_res)
|
|
print("Generating semantic segmentation without SAM")
|
|
if annotator_name == "seg_ufade20k":
|
|
original_semseg = _uniformer(input_image)
|
|
print("Generating semantic segmentation with SAM")
|
|
import annotator.uniformer as uniformer
|
|
original_uniformer_inference_segmentor = uniformer.inference_segmentor
|
|
uniformer.inference_segmentor = inject_uniformer_inference_segmentor
|
|
sam_semseg = _uniformer(input_image)
|
|
uniformer.inference_segmentor = original_uniformer_inference_segmentor
|
|
original_semseg = remove_pad(original_semseg)
|
|
sam_semseg = remove_pad(sam_semseg)
|
|
input_image = remove_pad(input_image)
|
|
output_gallery = [original_semseg, sam_semseg, blend_image_and_seg(input_image, original_semseg), blend_image_and_seg(input_image, sam_semseg)]
|
|
return output_gallery, "Uniformer semantic segmentation of ade20k done. Left is segmentation before SAM, right is segmentation after SAM."
|
|
else:
|
|
dataset = annotator_name.split('_')[-1][2:]
|
|
original_semseg = _oneformer(input_image, dataset=dataset)
|
|
print("Generating semantic segmentation with SAM")
|
|
from annotator.oneformer import OneformerDetector
|
|
original_oneformer_call = OneformerDetector.__call__
|
|
OneformerDetector.__call__ = inject_oneformer_detector_call
|
|
sam_semseg = _oneformer(input_image, dataset=dataset)
|
|
OneformerDetector.__call__ = original_oneformer_call
|
|
original_semseg = remove_pad(original_semseg)
|
|
sam_semseg = remove_pad(sam_semseg)
|
|
input_image = remove_pad(input_image)
|
|
output_gallery = [original_semseg, sam_semseg, blend_image_and_seg(input_image, original_semseg), blend_image_and_seg(input_image, sam_semseg)]
|
|
return output_gallery, f"Oneformer semantic segmentation of {dataset} done. Left is segmentation before SAM, right is segmentation after SAM."
|
|
else:
|
|
return random_segmentation(input_image)
|
|
|
|
|
|
def categorical_mask_image(crop_processor, crop_processor_res, crop_category_input, crop_input_image,
|
|
crop_pixel_perfect, crop_resize_mode, target_W, target_H):
|
|
if crop_input_image is None:
|
|
return "No input image."
|
|
if not os.path.isdir(os.path.join(scripts.basedir(), "annotator")) and not create_symbolic_link():
|
|
return "ControlNet extension not found."
|
|
filter_classes = crop_category_input.split('+')
|
|
if len(filter_classes) == 0:
|
|
return "No class selected."
|
|
try:
|
|
filter_classes = [int(i) for i in filter_classes]
|
|
except:
|
|
return "Illegal class id. You may have input some string."
|
|
crop_input_image_np = np.array(crop_input_image)
|
|
crop_processor_res = pixel_perfect_lllyasviel(crop_input_image_np, crop_processor_res, crop_pixel_perfect, crop_resize_mode, target_W, target_H)
|
|
crop_input_image, remove_pad = resize_image_with_pad(crop_input_image_np, crop_processor_res)
|
|
crop_input_image_copy = copy.deepcopy(crop_input_image)
|
|
global original_uniformer_inference_segmentor
|
|
print(f"Generating categories with processor {crop_processor}")
|
|
if crop_processor == "seg_ufade20k":
|
|
import annotator.uniformer as uniformer
|
|
original_uniformer_inference_segmentor = uniformer.inference_segmentor
|
|
uniformer.inference_segmentor = inject_uniformer_inference_segmentor
|
|
tmp_ouis = uniformer.show_result_pyplot
|
|
uniformer.show_result_pyplot = inject_uniformer_show_result_pyplot
|
|
sam_semseg = _uniformer(crop_input_image)
|
|
uniformer.inference_segmentor = original_uniformer_inference_segmentor
|
|
uniformer.show_result_pyplot = tmp_ouis
|
|
else:
|
|
dataset = crop_processor.split('_')[-1][2:]
|
|
from annotator.oneformer import OneformerDetector
|
|
original_oneformer_call = OneformerDetector.__call__
|
|
OneformerDetector.__call__ = inject_oneformer_detector_call_categorical_mask
|
|
sam_semseg = _oneformer(crop_input_image, dataset=dataset)
|
|
OneformerDetector.__call__ = original_oneformer_call
|
|
sam_semseg = remove_pad(sam_semseg)
|
|
mask = np.zeros(sam_semseg.shape, dtype=np.bool_)
|
|
for i in filter_classes:
|
|
mask[sam_semseg == i] = True
|
|
return mask, remove_pad(crop_input_image_copy)
|
|
|
|
|
|
def register_auto_sam(sam,
|
|
auto_sam_points_per_side, auto_sam_points_per_batch, auto_sam_pred_iou_thresh,
|
|
auto_sam_stability_score_thresh, auto_sam_stability_score_offset, auto_sam_box_nms_thresh,
|
|
auto_sam_crop_n_layers, auto_sam_crop_nms_thresh, auto_sam_crop_overlap_ratio,
|
|
auto_sam_crop_n_points_downscale_factor, auto_sam_min_mask_region_area, auto_sam_output_mode):
|
|
global global_sam
|
|
global_sam = SamAutomaticMaskGenerator(
|
|
sam, auto_sam_points_per_side, auto_sam_points_per_batch,
|
|
auto_sam_pred_iou_thresh, auto_sam_stability_score_thresh,
|
|
auto_sam_stability_score_offset, auto_sam_box_nms_thresh,
|
|
auto_sam_crop_n_layers, auto_sam_crop_nms_thresh, auto_sam_crop_overlap_ratio,
|
|
auto_sam_crop_n_points_downscale_factor, None,
|
|
auto_sam_min_mask_region_area, auto_sam_output_mode)
|
|
|
|
|
|
def pixel_perfect_lllyasviel(input_image, processor_res, use_pixel_perfect, resize_mode, target_W, target_H):
|
|
preprocessor_resolution = processor_res
|
|
if use_pixel_perfect:
|
|
raw_H, raw_W, _ = input_image.shape
|
|
k0 = float(target_H) / float(raw_H)
|
|
k1 = float(target_W) / float(raw_W)
|
|
if resize_mode == 2:
|
|
estimation = min(k0, k1) * float(min(raw_H, raw_W))
|
|
else:
|
|
estimation = max(k0, k1) * float(min(raw_H, raw_W))
|
|
preprocessor_resolution = int(np.round(float(estimation) / 64.0)) * 64
|
|
print(f'Pixel Perfect Mode (written by lllyasviel) Enabled.')
|
|
print(f'resize_mode = {str(resize_mode)}')
|
|
print(f'raw_H = {raw_H}')
|
|
print(f'raw_W = {raw_W}')
|
|
print(f'target_H = {target_H}')
|
|
print(f'target_W = {target_W}')
|
|
print(f'estimation = {estimation}')
|
|
print(f'preprocessor resolution = {preprocessor_resolution}')
|
|
return preprocessor_resolution
|