extensions/dddetailer/scripts/dddetailer.py

import os
import sys
from copy import copy
from pathlib import Path
from textwrap import dedent

import cv2
import gradio as gr
import numpy as np
from basicsr.utils.download_util import load_file_from_url
from packaging.version import parse
from PIL import Image

from launch import run
from modules import (
    devices,
    images,
    modelloader,
    processing,
    script_callbacks,
    scripts,
    shared,
)
from modules.paths import data_path, models_path
from modules.processing import (
    Processed,
    StableDiffusionProcessingImg2Img,
    StableDiffusionProcessingTxt2Img,
)
from modules.sd_models import model_hash
from modules.shared import cmd_opts, opts, state

DETECTION_DETAILER = "Detection Detailer"
dd_models_path = os.path.join(models_path, "mmdet")
python = sys.executable


def check_ddetailer() -> bool:
    try:
        from modules.paths import extensions_dir

        extensions_path = Path(extensions_dir)
    except ImportError:
        from modules.paths import data_path

        extensions_path = Path(data_path, "extensions")

    ddetailer_exists = any(p.is_dir() and p.name.startswith("ddetailer") for p in extensions_path.iterdir())
    return not ddetailer_exists


def check_install() -> bool:
    try:
        import mmcv
        import mmdet
        from mmdet.evaluation import get_classes
    except Exception:
        return False

    if not hasattr(mmcv, "__version__") or not hasattr(mmdet, "__version__"):
        return False

    v1 = parse(mmcv.__version__) >= parse("2.0.0")
    v2 = parse(mmdet.__version__) >= parse("3.0.0")
    return v1 and v2


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(("\\", "/")):
            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():
    if not check_ddetailer():
        message = """
        [-] dddetailer: dddetailer doesn't work with the original ddetailer extension.
                        dddetailer는 원본 ddetailer 확장이 있을 때 동작하지 않습니다.
        """
        raise RuntimeError(dedent(message))

    if not check_install():
        run(f'"{python}" -m pip uninstall -y mmcv mmcv-full mmdet mmengine')
        run(f'"{python}" -m pip install openmim', desc="Installing openmim", errdesc="Couldn't install openmim")
        run(
            f'"{python}" -m mim install mmcv>=2.0.0 mmdet>=3.0.0',
            desc="Installing mmdet",
            errdesc="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")
        # bbox
        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://raw.githubusercontent.com/Bing-su/dddetailer/master/config/mmdet_anime-face_yolov3.py",
            bbox_path,
        )
        # segm
        load_file_from_url(
            "https://github.com/Bing-su/dddetailer/releases/download/segm/mmdet_dd-person_mask2former.pth",
            segm_path,
        )
        load_file_from_url(
            "https://raw.githubusercontent.com/Bing-su/dddetailer/master/config/mmdet_dd-person_mask2former.py",
            segm_path,
        )
        load_file_from_url(
            "https://raw.githubusercontent.com/Bing-su/dddetailer/master/config/mask2former_r50_8xb2-lsj-50e_coco-panoptic.py",
            segm_path,
        )
        load_file_from_url(
            "https://raw.githubusercontent.com/Bing-su/dddetailer/master/config/coco_panoptic.py",
            segm_path,
        )


startup()


def gr_show(visible=True):
    return {"visible": visible, "__type__": "update"}


def ddetailer_extra_generation_params(
    dd_prompt,
    dd_neg_prompt,
    dd_model_a,
    dd_conf_a,
    dd_dilation_factor_a,
    dd_offset_x_a,
    dd_offset_y_a,
    dd_preprocess_b,
    dd_bitwise_op,
    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_cfg_scale,
):
    params = {
        "DDetailer prompt": dd_prompt,
        "DDetailer neg prompt": dd_neg_prompt,
        "DDetailer model a": dd_model_a,
        "DDetailer conf a": dd_conf_a,
        "DDetailer dilation a": dd_dilation_factor_a,
        "DDetailer offset x a": dd_offset_x_a,
        "DDetailer offset y a": dd_offset_y_a,
        "DDetailer preprocess b": dd_preprocess_b,
        "DDetailer bitwise": dd_bitwise_op,
        "DDetailer model b": dd_model_b,
        "DDetailer conf b": dd_conf_b,
        "DDetailer dilation b": dd_dilation_factor_b,
        "DDetailer offset x b": dd_offset_x_b,
        "DDetailer offset y b": dd_offset_y_b,
        "DDetailer mask blur": dd_mask_blur,
        "DDetailer denoising": dd_denoising_strength,
        "DDetailer inpaint full": dd_inpaint_full_res,
        "DDetailer inpaint padding": dd_inpaint_full_res_padding,
        "DDetailer cfg": dd_cfg_scale,
        "Script": DETECTION_DETAILER,
    }
    if not dd_prompt:
        params.pop("DDetailer prompt")
    if not dd_neg_prompt:
        params.pop("DDetailer neg prompt")
    return params


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(
                '<p style="margin-bottom:0.75em">Recommended settings: Use from inpaint tab, inpaint at full res ON, denoise < 0.5</p>'
            )
        else:
            info = gr.HTML("")
        dd_prompt = None
        with gr.Group():
            if not is_img2img:
                with gr.Row():
                    dd_prompt = gr.Textbox(
                        label="dd_prompt",
                        elem_id="t2i_dd_prompt",
                        show_label=False,
                        lines=3,
                        placeholder="Ddetailer Prompt",
                    )

                with gr.Row():
                    dd_neg_prompt = gr.Textbox(
                        label="dd_neg_prompt",
                        elem_id="t2i_dd_neg_prompt",
                        show_label=False,
                        lines=2,
                        placeholder="Ddetailer Negative prompt",
                    )

            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=True,
                )
                dd_dilation_factor_a = gr.Slider(
                    label="Dilation factor (A)",
                    minimum=0,
                    maximum=255,
                    step=1,
                    value=4,
                    visible=True,
                )

            with gr.Row():
                dd_offset_x_a = gr.Slider(
                    label="X offset (A)",
                    minimum=-200,
                    maximum=200,
                    step=1,
                    value=0,
                    visible=True,
                )
                dd_offset_y_a = gr.Slider(
                    label="Y offset (A)",
                    minimum=-200,
                    maximum=200,
                    step=1,
                    value=0,
                    visible=True,
                )

            with gr.Row():
                dd_preprocess_b = gr.Checkbox(
                    label="Inpaint model B detections before model A runs",
                    value=False,
                    visible=True,
                )
                dd_bitwise_op = gr.Radio(
                    label="Bitwise operation",
                    choices=["None", "A&B", "A-B"],
                    value="None",
                    visible=True,
                )

        br = gr.HTML("<br>")

        with gr.Group():
            with gr.Row():
                dd_model_b = gr.Dropdown(
                    label="Secondary detection model (B) (optional)",
                    choices=model_list,
                    value="None",
                    visible=True,
                    type="value",
                )

            with gr.Row():
                dd_conf_b = gr.Slider(
                    label="Detection confidence threshold % (B)",
                    minimum=0,
                    maximum=100,
                    step=1,
                    value=30,
                    visible=True,
                )
                dd_dilation_factor_b = gr.Slider(
                    label="Dilation factor (B)",
                    minimum=0,
                    maximum=255,
                    step=1,
                    value=4,
                    visible=True,
                )

            with gr.Row():
                dd_offset_x_b = gr.Slider(
                    label="X offset (B)",
                    minimum=-200,
                    maximum=200,
                    step=1,
                    value=0,
                    visible=True,
                )
                dd_offset_y_b = gr.Slider(
                    label="Y offset (B)",
                    minimum=-200,
                    maximum=200,
                    step=1,
                    value=0,
                    visible=True,
                )

        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_cfg_scale = gr.Slider(
                    label="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,
            ],
        )
        if dd_prompt:
            self.infotext_fields = (
                (dd_prompt, "DDetailer prompt"),
                (dd_neg_prompt, "DDetailer neg prompt"),
                (dd_model_a, "DDetailer model a"),
                (dd_conf_a, "DDetailer conf a"),
                (dd_dilation_factor_a, "DDetailer dilation a"),
                (dd_offset_x_a, "DDetailer offset x a"),
                (dd_offset_y_a, "DDetailer offset y a"),
                (dd_preprocess_b, "DDetailer preprocess b"),
                (dd_bitwise_op, "DDetailer bitwise"),
                (dd_model_b, "DDetailer model b"),
                (dd_conf_b, "DDetailer conf b"),
                (dd_dilation_factor_b, "DDetailer dilation b"),
                (dd_offset_x_b, "DDetailer offset x b"),
                (dd_offset_y_b, "DDetailer offset y b"),
                (dd_mask_blur, "DDetailer mask blur"),
                (dd_denoising_strength, "DDetailer denoising"),
                (dd_inpaint_full_res, "DDetailer inpaint full"),
                (dd_inpaint_full_res_padding, "DDetailer inpaint padding"),
                (dd_cfg_scale, "DDetailer cfg"),
            )

        ret = [
            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_cfg_scale,
        ]
        if not is_img2img:
            ret += [dd_prompt, dd_neg_prompt]
        return ret

    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_cfg_scale,
        dd_prompt=None,
        dd_neg_prompt=None,
    ):
        processing.fix_seed(p)
        seed = p.seed
        subseed = p.subseed
        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)
        info = ""

        # ddetailer info
        extra_generation_params = ddetailer_extra_generation_params(
            dd_prompt,
            dd_neg_prompt,
            dd_model_a,
            dd_conf_a,
            dd_dilation_factor_a,
            dd_offset_x_a,
            dd_offset_y_a,
            dd_preprocess_b,
            dd_bitwise_op,
            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_cfg_scale,
        )
        p.extra_generation_params.update(extra_generation_params)

        p_txt = copy(p)
        if not is_txt2img:
            orig_image = p.init_images[0]
        else:
            img2img_sampler_name = p_txt.sampler_name
            # PLMS/UniPC do not support img2img so we just silently switch to DDIM
            if p_txt.sampler_name in ["PLMS", "UniPC"]:
                img2img_sampler_name = "DDIM"
            p_txt_prompt = dd_prompt if dd_prompt else p_txt.prompt
            p_txt_neg_prompt = dd_neg_prompt if dd_neg_prompt else p_txt.negative_prompt
            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_neg_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=img2img_sampler_name,
                n_iter=p_txt.n_iter,
                steps=p_txt.steps,
                cfg_scale=p_txt.cfg_scale,
                width=p_txt.width,
                height=p_txt.height,
                tiling=p_txt.tiling,
                extra_generation_params=p_txt.extra_generation_params,
            )
            p.do_not_save_grid = True
            p.do_not_save_samples = True
            p.cached_c = [None, None]
            p.cached_uc = [None, None]

            p.scripts = p_txt.scripts
            p.script_args = p_txt.script_args

        # output info
        all_prompts = []
        all_negative_prompts = []
        all_seeds = []
        all_subseeds = []
        infotexts = []
        output_images = []

        state.job_count = ddetail_count
        for n in range(ddetail_count):
            devices.torch_gc()
            start_seed = seed + n

            all_prompts.append(p_txt.prompt)
            all_negative_prompts.append(p_txt.negative_prompt)
            all_seeds.append(start_seed)
            all_subseeds.append(subseed + 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]
                info = processed.info
                if not dd_prompt:
                    p.prompt = processed.all_prompts[0]
                if not dd_neg_prompt:
                    p.negative_prompt = processed.all_negative_prompts[0]
                all_prompts[n] = processed.all_prompts[0]
                all_negative_prompts[n] = processed.all_negative_prompts[0]
            else:
                init_image = orig_image
                p.prompt = p_txt.prompt
                p.negative_prompt = p_txt.negative_prompt
            p.cfg_scale = dd_cfg_scale

            if opts.enable_pnginfo:
                init_image.info["parameters"] = info

            infotexts.append(info)
            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)
                        if not is_txt2img:
                            p.prompt = processed.all_prompts[0]
                            p.negative_prompt = processed.all_negative_prompts[0]
                        p.seed = processed.seed + 1
                        p.subseed = processed.subseed + 1
                        p.init_images = [processed.images[0]]

                    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 not is_txt2img:
                            p.prompt = processed.all_prompts[0]
                            p.negative_prompt = processed.all_negative_prompts[0]
                            info = processed.info
                            all_prompts[n] = processed.all_prompts[0]
                            all_negative_prompts[n] = processed.all_negative_prompts[0]
                        p.seed = processed.seed + 1
                        p.subseed = processed.subseed + 1
                        p.init_images = [processed.images[0]]

                    if gen_count > 0:
                        final_image = processed.images[0]

                        if opts.enable_pnginfo:
                            final_image.info["parameters"] = info
                        output_images[n] = final_image
                        infotexts[n] = info

                        if opts.samples_save:
                            images.save_image(
                                final_image,
                                p.outpath_samples,
                                "",
                                start_seed,
                                p.prompt,
                                opts.samples_format,
                                info=info,
                                p=p,
                            )

                else:
                    print(f"No model {label_a} detections for output generation {n} with current settings.")

                    if opts.samples_save:
                        images.save_image(
                            init_image,
                            p.outpath_samples,
                            "",
                            start_seed,
                            p.prompt,
                            opts.samples_format,
                            info=info,
                            p=p,
                        )

            state.job = f"Generation {n + 1} out of {state.job_count}"

        if dd_prompt or dd_neg_prompt:
            params_txt = os.path.join(data_path, "params.txt")
            with open(params_txt, "w", encoding="utf-8") as file:
                file.write(infotexts[0])

        return Processed(
            p,
            output_images,
            seed,
            infotexts[0],
            all_prompts=all_prompts,
            all_negative_prompts=all_negative_prompts,
            all_seeds=all_seeds,
            all_subseeds=all_subseeds,
            infotexts=infotexts,
        )


def modeldataset(model_shortname):
    path = modelpath(model_shortname)
    dataset = "coco" if "mmdet" in path and "segm" in path else "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
    return None


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]
    scores = results[3]

    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 = scores[i]
        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/dddetailer/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/dddetailer/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


from mmdet.apis import inference_detector, init_detector
from mmdet.evaluation import get_classes


def get_device():
    device = devices.get_optimal_device_name()
    if device == "mps":
        return device
    if any(getattr(cmd_opts, vram, False) for vram in ["lowvram", "medvram"]):
        return "cpu"
    return 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)).pred_instances
    bboxes = mmdet_results.bboxes.cpu().numpy()
    segms = mmdet_results.masks.cpu().numpy()
    scores = mmdet_results.scores.cpu().numpy()
    dataset = modeldataset(modelname)
    classes = get_classes(dataset)

    n, m = bboxes.shape
    if n == 0:
        return [[], [], [], []]
    labels = mmdet_results.labels
    filter_inds = np.where(scores > 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])
        results[3].append(scores[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)
    output = inference_detector(model, np.array(image)).pred_instances
    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 in output.bboxes:
        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 = output.bboxes.shape
    if n == 0:
        return [[], [], [], []]
    bboxes = output.bboxes.cpu().numpy()
    scores = output.scores.cpu().numpy()
    filter_inds = np.where(scores > conf_thres)[0]
    results = [[], [], [], []]
    for i in filter_inds:
        results[0].append(label)
        results[1].append(bboxes[i])
        results[2].append(segms[i])
        results[3].append(scores[i])

    return results


script_callbacks.on_ui_settings(on_ui_settings)