import os import sys import cv2 from PIL import Image import numpy as np import gradio as gr from modules import processing, images from modules import scripts, script_callbacks, shared, devices, modelloader from modules.processing import Processed, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img from modules.shared import opts, cmd_opts, state from modules.sd_models import model_hash from modules.paths import models_path from basicsr.utils.download_util import load_file_from_url dd_models_path = os.path.join(models_path, "mmdet") def list_models(model_path): model_list = modelloader.load_models(model_path=model_path, ext_filter=[".pth"]) def modeltitle(path, shorthash): abspath = os.path.abspath(path) if abspath.startswith(model_path): name = abspath.replace(model_path, '') else: name = os.path.basename(path) if name.startswith("\\") or name.startswith("/"): name = name[1:] shortname = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0] return f'{name} [{shorthash}]', shortname models = [] for filename in model_list: h = model_hash(filename) title, short_model_name = modeltitle(filename, h) models.append(title) return models def startup(): from launch import is_installed, run if not is_installed("mmdet"): python = sys.executable run(f'"{python}" -m pip install -U openmim==0.3.7', desc="Installing openmim", errdesc="Couldn't install openmim") run(f'"{python}" -m mim install mmcv-full==1.7.1', desc=f"Installing mmcv-full", errdesc=f"Couldn't install mmcv-full") run(f'"{python}" -m pip install mmdet==2.28.2', desc=f"Installing mmdet", errdesc=f"Couldn't install mmdet") if (len(list_models(dd_models_path)) == 0): print("No detection models found, downloading...") bbox_path = os.path.join(dd_models_path, "bbox") segm_path = os.path.join(dd_models_path, "segm") load_file_from_url("https://huggingface.co/dustysys/ddetailer/resolve/main/mmdet/bbox/mmdet_anime-face_yolov3.pth", bbox_path) load_file_from_url("https://huggingface.co/dustysys/ddetailer/raw/main/mmdet/bbox/mmdet_anime-face_yolov3.py", bbox_path) load_file_from_url("https://huggingface.co/dustysys/ddetailer/resolve/main/mmdet/segm/mmdet_dd-person_mask2former.pth", segm_path) load_file_from_url("https://huggingface.co/dustysys/ddetailer/raw/main/mmdet/segm/mmdet_dd-person_mask2former.py", segm_path) startup() def gr_show(visible=True): return {"visible": visible, "__type__": "update"} class DetectionDetailerScript(scripts.Script): def title(self): return "Detection Detailer" def show(self, is_img2img): return True def ui(self, is_img2img): import modules.ui model_list = list_models(dd_models_path) model_list.insert(0, "None") if is_img2img: info = gr.HTML("

Recommended settings: Use from inpaint tab, inpaint at full res ON, denoise <0.5

") else: info = gr.HTML("") with gr.Group(): with gr.Row(): dd_model_a = gr.Dropdown(label="Primary detection model (A)", choices=model_list,value = "None", visible=True, type="value") with gr.Row(): dd_conf_a = gr.Slider(label='Detection confidence threshold % (A)', minimum=0, maximum=100, step=1, value=30, visible=False) dd_dilation_factor_a = gr.Slider(label='Dilation factor (A)', minimum=0, maximum=255, step=1, value=4, visible=False) with gr.Row(): dd_offset_x_a = gr.Slider(label='X offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False) dd_offset_y_a = gr.Slider(label='Y offset (A)', minimum=-200, maximum=200, step=1, value=0, visible=False) with gr.Row(): dd_preprocess_b = gr.Checkbox(label='Inpaint model B detections before model A runs', value=False, visible=False) dd_bitwise_op = gr.Radio(label='Bitwise operation', choices=['None', 'A&B', 'A-B'], value="None", visible=False) br = gr.HTML("
") with gr.Group(): with gr.Row(): dd_model_b = gr.Dropdown(label="Secondary detection model (B) (optional)", choices=model_list,value = "None", visible =False, type="value") with gr.Row(): dd_conf_b = gr.Slider(label='Detection confidence threshold % (B)', minimum=0, maximum=100, step=1, value=30, visible=False) dd_dilation_factor_b = gr.Slider(label='Dilation factor (B)', minimum=0, maximum=255, step=1, value=4, visible=False) with gr.Row(): dd_offset_x_b = gr.Slider(label='X offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False) dd_offset_y_b = gr.Slider(label='Y offset (B)', minimum=-200, maximum=200, step=1, value=0, visible=False) with gr.Group(): with gr.Row(): dd_mask_blur = gr.Slider(label='Mask blur ', minimum=0, maximum=64, step=1, value=4, visible=(not is_img2img)) dd_denoising_strength = gr.Slider(label='Denoising strength (Inpaint)', minimum=0.0, maximum=1.0, step=0.01, value=0.4, visible=(not is_img2img)) with gr.Row(): dd_inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution ', value=True, visible = (not is_img2img)) dd_inpaint_full_res_padding = gr.Slider(label='Inpaint at full resolution padding, pixels ', minimum=0, maximum=256, step=4, value=32, visible=(not is_img2img)) with gr.Row(): dd_mimic_cfg = gr.Slider(label='Mimic CFG Scale', minimum=0, maximum=30, step=0.5, value=7, visible=True) dd_model_a.change( lambda modelname: { dd_model_b:gr_show( modelname != "None" ), dd_conf_a:gr_show( modelname != "None" ), dd_dilation_factor_a:gr_show( modelname != "None"), dd_offset_x_a:gr_show( modelname != "None" ), dd_offset_y_a:gr_show( modelname != "None" ) }, inputs= [dd_model_a], outputs =[dd_model_b, dd_conf_a, dd_dilation_factor_a, dd_offset_x_a, dd_offset_y_a] ) dd_model_b.change( lambda modelname: { dd_preprocess_b:gr_show( modelname != "None" ), dd_bitwise_op:gr_show( modelname != "None" ), dd_conf_b:gr_show( modelname != "None" ), dd_dilation_factor_b:gr_show( modelname != "None"), dd_offset_x_b:gr_show( modelname != "None" ), dd_offset_y_b:gr_show( modelname != "None" ) }, inputs= [dd_model_b], outputs =[dd_preprocess_b, dd_bitwise_op, dd_conf_b, dd_dilation_factor_b, dd_offset_x_b, dd_offset_y_b] ) return [info, dd_model_a, dd_conf_a, dd_dilation_factor_a, dd_offset_x_a, dd_offset_y_a, dd_preprocess_b, dd_bitwise_op, br, dd_model_b, dd_conf_b, dd_dilation_factor_b, dd_offset_x_b, dd_offset_y_b, dd_mask_blur, dd_denoising_strength, dd_inpaint_full_res, dd_inpaint_full_res_padding, dd_mimic_cfg ] def run(self, p, info, dd_model_a, dd_conf_a, dd_dilation_factor_a, dd_offset_x_a, dd_offset_y_a, dd_preprocess_b, dd_bitwise_op, br, dd_model_b, dd_conf_b, dd_dilation_factor_b, dd_offset_x_b, dd_offset_y_b, dd_mask_blur, dd_denoising_strength, dd_inpaint_full_res, dd_inpaint_full_res_padding, dd_mimic_cfg): processing.fix_seed(p) initial_info = None seed = p.seed p.batch_size = 1 ddetail_count = p.n_iter p.n_iter = 1 p.do_not_save_grid = True p.do_not_save_samples = True is_txt2img = isinstance(p, StableDiffusionProcessingTxt2Img) if (not is_txt2img): orig_image = p.init_images[0] else: p_txt = p print(f"mimic_scale = {dd_mimic_cfg}") p = StableDiffusionProcessingImg2Img( init_images = None, resize_mode = 0, denoising_strength = dd_denoising_strength, mask = None, mask_blur= dd_mask_blur, inpainting_fill = 1, inpaint_full_res = dd_inpaint_full_res, inpaint_full_res_padding= dd_inpaint_full_res_padding, inpainting_mask_invert= 0, sd_model=p_txt.sd_model, outpath_samples=p_txt.outpath_samples, outpath_grids=p_txt.outpath_grids, prompt=p_txt.prompt, negative_prompt=p_txt.negative_prompt, styles=p_txt.styles, seed=p_txt.seed, subseed=p_txt.subseed, subseed_strength=p_txt.subseed_strength, seed_resize_from_h=p_txt.seed_resize_from_h, seed_resize_from_w=p_txt.seed_resize_from_w, sampler_name=p_txt.sampler_name, n_iter=p_txt.n_iter, steps=p_txt.steps, cfg_scale=dd_mimic_cfg, width=p_txt.width, height=p_txt.height, tiling=p_txt.tiling, ) p.do_not_save_grid = True p.do_not_save_samples = True output_images = [] state.job_count = ddetail_count for n in range(ddetail_count): devices.torch_gc() start_seed = seed + n if ( is_txt2img ): print(f"Processing initial image for output generation {n + 1}.") p_txt.seed = start_seed processed = processing.process_images(p_txt) init_image = processed.images[0] else: init_image = orig_image output_images.append(init_image) masks_a = [] masks_b_pre = [] # Optional secondary pre-processing run if (dd_model_b != "None" and dd_preprocess_b): label_b_pre = "B" results_b_pre = inference(init_image, dd_model_b, dd_conf_b/100.0, label_b_pre) masks_b_pre = create_segmasks(results_b_pre) masks_b_pre = dilate_masks(masks_b_pre, dd_dilation_factor_b, 1) masks_b_pre = offset_masks(masks_b_pre,dd_offset_x_b, dd_offset_y_b) if (len(masks_b_pre) > 0): results_b_pre = update_result_masks(results_b_pre, masks_b_pre) segmask_preview_b = create_segmask_preview(results_b_pre, init_image) shared.state.current_image = segmask_preview_b if ( opts.dd_save_previews): images.save_image(segmask_preview_b, opts.outdir_ddetailer_previews, "", start_seed, p.prompt, opts.samples_format, p=p) gen_count = len(masks_b_pre) state.job_count += gen_count print(f"Processing {gen_count} model {label_b_pre} detections for output generation {n + 1}.") p.seed = start_seed p.init_images = [init_image] for i in range(gen_count): p.image_mask = masks_b_pre[i] if ( opts.dd_save_masks): images.save_image(masks_b_pre[i], opts.outdir_ddetailer_masks, "", start_seed, p.prompt, opts.samples_format, p=p) processed = processing.process_images(p) p.seed = processed.seed + 1 p.init_images = processed.images if (gen_count > 0): output_images[n] = processed.images[0] init_image = processed.images[0] else: print(f"No model B detections for output generation {n} with current settings.") # Primary run if (dd_model_a != "None"): label_a = "A" if (dd_model_b != "None" and dd_bitwise_op != "None"): label_a = dd_bitwise_op results_a = inference(init_image, dd_model_a, dd_conf_a/100.0, label_a) masks_a = create_segmasks(results_a) masks_a = dilate_masks(masks_a, dd_dilation_factor_a, 1) masks_a = offset_masks(masks_a,dd_offset_x_a, dd_offset_y_a) if (dd_model_b != "None" and dd_bitwise_op != "None"): label_b = "B" results_b = inference(init_image, dd_model_b, dd_conf_b/100.0, label_b) masks_b = create_segmasks(results_b) masks_b = dilate_masks(masks_b, dd_dilation_factor_b, 1) masks_b = offset_masks(masks_b,dd_offset_x_b, dd_offset_y_b) if (len(masks_b) > 0): combined_mask_b = combine_masks(masks_b) for i in reversed(range(len(masks_a))): if (dd_bitwise_op == "A&B"): masks_a[i] = bitwise_and_masks(masks_a[i], combined_mask_b) elif (dd_bitwise_op == "A-B"): masks_a[i] = subtract_masks(masks_a[i], combined_mask_b) if (is_allblack(masks_a[i])): del masks_a[i] for result in results_a: del result[i] else: print("No model B detections to overlap with model A masks") results_a = [] masks_a = [] if (len(masks_a) > 0): results_a = update_result_masks(results_a, masks_a) segmask_preview_a = create_segmask_preview(results_a, init_image) shared.state.current_image = segmask_preview_a if ( opts.dd_save_previews): images.save_image(segmask_preview_a, opts.outdir_ddetailer_previews, "", start_seed, p.prompt, opts.samples_format, p=p) gen_count = len(masks_a) state.job_count += gen_count print(f"Processing {gen_count} model {label_a} detections for output generation {n + 1}.") p.seed = start_seed p.init_images = [init_image] for i in range(gen_count): p.image_mask = masks_a[i] if ( opts.dd_save_masks): images.save_image(masks_a[i], opts.outdir_ddetailer_masks, "", start_seed, p.prompt, opts.samples_format, p=p) processed = processing.process_images(p) if initial_info is None: initial_info = processed.info p.seed = processed.seed + 1 p.init_images = processed.images if (gen_count > 0): output_images[n] = processed.images[0] if ( opts.samples_save ): images.save_image(processed.images[0], p.outpath_samples, "", start_seed, p.prompt, opts.samples_format, info=initial_info, p=p) else: print(f"No model {label_a} detections for output generation {n} with current settings.") state.job = f"Generation {n + 1} out of {state.job_count}" if (initial_info is None): initial_info = "No detections found." return Processed(p, output_images, seed, initial_info) def modeldataset(model_shortname): path = modelpath(model_shortname) if ("mmdet" in path and "segm" in path): dataset = 'coco' else: dataset = 'bbox' return dataset def modelpath(model_shortname): model_list = modelloader.load_models(model_path=dd_models_path, ext_filter=[".pth"]) model_h = model_shortname.split("[")[-1].split("]")[0] for path in model_list: if ( model_hash(path) == model_h): return path def update_result_masks(results, masks): for i in range(len(masks)): boolmask = np.array(masks[i], dtype=bool) results[2][i] = boolmask return results def create_segmask_preview(results, image): labels = results[0] bboxes = results[1] segms = results[2] cv2_image = np.array(image) cv2_image = cv2_image[:, :, ::-1].copy() for i in range(len(segms)): color = np.full_like(cv2_image, np.random.randint(100, 256, (1, 3), dtype=np.uint8)) alpha = 0.2 color_image = cv2.addWeighted(cv2_image, alpha, color, 1-alpha, 0) cv2_mask = segms[i].astype(np.uint8) * 255 cv2_mask_bool = np.array(segms[i], dtype=bool) centroid = np.mean(np.argwhere(cv2_mask_bool),axis=0) centroid_x, centroid_y = int(centroid[1]), int(centroid[0]) cv2_mask_rgb = cv2.merge((cv2_mask, cv2_mask, cv2_mask)) cv2_image = np.where(cv2_mask_rgb == 255, color_image, cv2_image) text_color = tuple([int(x) for x in ( color[0][0] - 100 )]) name = labels[i] score = bboxes[i][4] score = str(score)[:4] text = name + ":" + score cv2.putText(cv2_image, text, (centroid_x - 30, centroid_y), cv2.FONT_HERSHEY_DUPLEX, 0.4, text_color, 1, cv2.LINE_AA) if ( len(segms) > 0): preview_image = Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)) else: preview_image = image return preview_image def is_allblack(mask): cv2_mask = np.array(mask) return cv2.countNonZero(cv2_mask) == 0 def bitwise_and_masks(mask1, mask2): cv2_mask1 = np.array(mask1) cv2_mask2 = np.array(mask2) cv2_mask = cv2.bitwise_and(cv2_mask1, cv2_mask2) mask = Image.fromarray(cv2_mask) return mask def subtract_masks(mask1, mask2): cv2_mask1 = np.array(mask1) cv2_mask2 = np.array(mask2) cv2_mask = cv2.subtract(cv2_mask1, cv2_mask2) mask = Image.fromarray(cv2_mask) return mask def dilate_masks(masks, dilation_factor, iter=1): if dilation_factor == 0: return masks dilated_masks = [] kernel = np.ones((dilation_factor,dilation_factor), np.uint8) for i in range(len(masks)): cv2_mask = np.array(masks[i]) dilated_mask = cv2.dilate(cv2_mask, kernel, iter) dilated_masks.append(Image.fromarray(dilated_mask)) return dilated_masks def offset_masks(masks, offset_x, offset_y): if (offset_x == 0 and offset_y == 0): return masks offset_masks = [] for i in range(len(masks)): cv2_mask = np.array(masks[i]) offset_mask = cv2_mask.copy() offset_mask = np.roll(offset_mask, -offset_y, axis=0) offset_mask = np.roll(offset_mask, offset_x, axis=1) offset_masks.append(Image.fromarray(offset_mask)) return offset_masks def combine_masks(masks): initial_cv2_mask = np.array(masks[0]) combined_cv2_mask = initial_cv2_mask for i in range(1, len(masks)): cv2_mask = np.array(masks[i]) combined_cv2_mask = cv2.bitwise_or(combined_cv2_mask, cv2_mask) combined_mask = Image.fromarray(combined_cv2_mask) return combined_mask def on_ui_settings(): shared.opts.add_option("dd_save_previews", shared.OptionInfo(False, "Save mask previews", section=("ddetailer", "Detection Detailer"))) shared.opts.add_option("outdir_ddetailer_previews", shared.OptionInfo("extensions/ddetailer/outputs/masks-previews", 'Output directory for mask previews', section=("ddetailer", "Detection Detailer"))) shared.opts.add_option("dd_save_masks", shared.OptionInfo(False, "Save masks", section=("ddetailer", "Detection Detailer"))) shared.opts.add_option("outdir_ddetailer_masks", shared.OptionInfo("extensions/ddetailer/outputs/masks", 'Output directory for masks', section=("ddetailer", "Detection Detailer"))) def create_segmasks(results): segms = results[2] segmasks = [] for i in range(len(segms)): cv2_mask = segms[i].astype(np.uint8) * 255 mask = Image.fromarray(cv2_mask) segmasks.append(mask) return segmasks import mmcv from mmdet.core import get_classes from mmdet.apis import (inference_detector, init_detector) def get_device(): device_id = shared.cmd_opts.device_id if device_id is not None: cuda_device = f"cuda:{device_id}" else: cuda_device = "cpu" return cuda_device def inference(image, modelname, conf_thres, label): path = modelpath(modelname) if ( "mmdet" in path and "bbox" in path ): results = inference_mmdet_bbox(image, modelname, conf_thres, label) elif ( "mmdet" in path and "segm" in path): results = inference_mmdet_segm(image, modelname, conf_thres, label) return results def inference_mmdet_segm(image, modelname, conf_thres, label): model_checkpoint = modelpath(modelname) model_config = os.path.splitext(model_checkpoint)[0] + ".py" model_device = get_device() model = init_detector(model_config, model_checkpoint, device=model_device) mmdet_results = inference_detector(model, np.array(image)) bbox_results, segm_results = mmdet_results dataset = modeldataset(modelname) classes = get_classes(dataset) labels = [ np.full(bbox.shape[0], i, dtype=np.int32) for i, bbox in enumerate(bbox_results) ] n,m = bbox_results[0].shape if (n == 0): return [[],[],[]] labels = np.concatenate(labels) bboxes = np.vstack(bbox_results) segms = mmcv.concat_list(segm_results) filter_inds = np.where(bboxes[:,-1] > conf_thres)[0] results = [[],[],[]] for i in filter_inds: results[0].append(label + "-" + classes[labels[i]]) results[1].append(bboxes[i]) results[2].append(segms[i]) return results def inference_mmdet_bbox(image, modelname, conf_thres, label): model_checkpoint = modelpath(modelname) model_config = os.path.splitext(model_checkpoint)[0] + ".py" model_device = get_device() model = init_detector(model_config, model_checkpoint, device=model_device) results = inference_detector(model, np.array(image)) cv2_image = np.array(image) cv2_image = cv2_image[:, :, ::-1].copy() cv2_gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY) segms = [] for (x0, y0, x1, y1, conf) in results[0]: cv2_mask = np.zeros((cv2_gray.shape), np.uint8) cv2.rectangle(cv2_mask, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1) cv2_mask_bool = cv2_mask.astype(bool) segms.append(cv2_mask_bool) n,m = results[0].shape if (n == 0): return [[],[],[]] bboxes = np.vstack(results[0]) filter_inds = np.where(bboxes[:,-1] > conf_thres)[0] results = [[],[],[]] for i in filter_inds: results[0].append(label) results[1].append(bboxes[i]) results[2].append(segms[i]) return results script_callbacks.on_ui_settings(on_ui_settings)