Documentation

Class name: MiDaS_Background_Foreground_Removal
Category: WAS Suite/Image/AI
Output node: False
Repo Ref: https://github.com/WASasquatch/was-node-suite-comfyui

该节点使用MiDaS模型从图像中估计深度，并根据用户输入移除背景或前景，增强图像的视觉清晰度和焦点。

Input types

Required

image
- 源图像，从中估计深度并移除背景或前景。
- Comfy dtype: IMAGE
- Python dtype: PIL.Image
midas_model
- 指定用于深度估计的MiDaS模型的大小，较大的模型通常能提供更准确的结果，但需要更多的计算资源。
- Comfy dtype: COMBO[ ['DPT_Large', 'DPT_Hybrid', 'DPT_Small' ]]
- Python dtype: str
remove
- 指示应从图像中移除背景还是前景，影响输出的构图。
- Comfy dtype: COMBO[ ['background', 'foreground' ]]
- Python dtype: str

Optional

use_cpu
- 决定是否使用CPU进行处理。卸载到GPU通常更快，但需要兼容的硬件。
- Comfy dtype: COMBO[ ['false', 'true' ]]
- Python dtype: str
threshold
- 对深度图应用阈值，可以细化移除元素与图像其余部分的分离。
- Comfy dtype: COMBO[ ['false', 'true' ]]
- Python dtype: str
threshold_low
- 设置深度值阈值范围的下限，影响分割的精度。
- Comfy dtype: FLOAT
- Python dtype: float
threshold_mid
- 定义阈值范围的中点，进一步调整保留和移除的图像部分之间的过渡。
- Comfy dtype: FLOAT
- Python dtype: float
threshold_high
- 建立阈值范围的上限，这可以影响移除过程中图像最深处的处理方式。
- Comfy dtype: FLOAT
- Python dtype: float
smoothing
- 控制应用于深度图的平滑程度，这有助于减少最终输出中的噪声和伪影。
- Comfy dtype: FLOAT
- Python dtype: float
background_red
- 设置用于填充移除区域的背景颜色的红色分量，影响结果图像的视觉外观。
- Comfy dtype: INT
- Python dtype: int
background_green
- 确定背景颜色的绿色分量，有助于构成移除元素区域的整体色彩方案。
- Comfy dtype: INT
- Python dtype: int
background_blue
- 指定背景颜色的蓝色分量，用于在编辑图像中实现期望的色彩平衡。
- Comfy dtype: INT
- Python dtype: int

Output types

result_image
- 处理后的图像，已移除指定的背景或前景元素，可供进一步使用或分析。
- Comfy dtype: IMAGE
- Python dtype: PIL.Image
depth_map
- 从输入图像派生出的深度图，视觉上表示场景中的估计距离，并在内部用于移除过程。
- Comfy dtype: IMAGE
- Python dtype: PIL.Image

Usage tips

Infra type: GPU

Source code

class MiDaS_Background_Foreground_Removal:

    def __init__(self):
        self.midas_dir = os.path.join(MODELS_DIR, 'midas')

    @classmethod
    def INPUT_TYPES(cls):
        return {'required': {'image': ('IMAGE',), 'use_cpu': (['false', 'true'],), 'midas_model': (['DPT_Large', 'DPT_Hybrid', 'DPT_Small'],), 'remove': (['background', 'foregroud'],), 'threshold': (['false', 'true'],), 'threshold_low': ('FLOAT', {'default': 10, 'min': 0, 'max': 255, 'step': 1}), 'threshold_mid': ('FLOAT', {'default': 200, 'min': 0, 'max': 255, 'step': 1}), 'threshold_high': ('FLOAT', {'default': 210, 'min': 0, 'max': 255, 'step': 1}), 'smoothing': ('FLOAT', {'default': 0.25, 'min': 0.0, 'max': 16.0, 'step': 0.01}), 'background_red': ('INT', {'default': 0, 'min': 0, 'max': 255, 'step': 1}), 'background_green': ('INT', {'default': 0, 'min': 0, 'max': 255, 'step': 1}), 'background_blue': ('INT', {'default': 0, 'min': 0, 'max': 255, 'step': 1})}}
    RETURN_TYPES = ('IMAGE', 'IMAGE')
    FUNCTION = 'midas_remove'
    CATEGORY = 'WAS Suite/Image/AI'

    def midas_remove(self, image, midas_model, use_cpu='false', remove='background', threshold='false', threshold_low=0, threshold_mid=127, threshold_high=255, smoothing=0.25, background_red=0, background_green=0, background_blue=0):
        global MIDAS_INSTALLED
        if not MIDAS_INSTALLED:
            self.install_midas()
        import cv2 as cv
        i = 255.0 * image.cpu().numpy().squeeze()
        img = i
        img_original = tensor2pil(image).convert('RGB')
        cstr('Downloading and loading MiDaS Model...').msg.print()
        torch.hub.set_dir(self.midas_dir)
        midas = torch.hub.load('intel-isl/MiDaS', midas_model, trust_repo=True)
        device = torch.device('cuda') if torch.cuda.is_available() and use_cpu == 'false' else torch.device('cpu')
        cstr(f'MiDaS is using device: {device}').msg.print()
        midas.to(device).eval()
        midas_transforms = torch.hub.load('intel-isl/MiDaS', 'transforms')
        if midas_model == 'DPT_Large' or midas_model == 'DPT_Hybrid':
            transform = midas_transforms.dpt_transform
        else:
            transform = midas_transforms.small_transform
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        input_batch = transform(img).to(device)
        cstr('Approximating depth from image.').msg.print()
        with torch.no_grad():
            prediction = midas(input_batch)
            prediction = torch.nn.functional.interpolate(prediction.unsqueeze(1), size=img.shape[:2], mode='bicubic', align_corners=False).squeeze()
        if remove == 'foreground':
            depth = 255 - prediction.cpu().numpy().astype(np.uint8)
            depth = depth.astype(np.float32)
        else:
            depth = prediction.cpu().numpy().astype(np.float32)
        depth = depth * 255 / np.max(depth) / 255
        depth = Image.fromarray(np.uint8(depth * 255))
        if threshold == 'true':
            levels = self.AdjustLevels(threshold_low, threshold_mid, threshold_high)
            depth = levels.adjust(depth.convert('RGB')).convert('L')
        if smoothing > 0:
            depth = depth.filter(ImageFilter.GaussianBlur(radius=smoothing))
        depth = depth.resize(img_original.size).convert('L')
        background_red = int(background_red) if isinstance(background_red, (int, float)) else 0
        background_green = int(background_green) if isinstance(background_green, (int, float)) else 0
        background_blue = int(background_blue) if isinstance(background_blue, (int, float)) else 0
        background_color = (background_red, background_green, background_blue)
        background = Image.new(mode='RGB', size=img_original.size, color=background_color)
        result_img = Image.composite(img_original, background, depth)
        del midas, device, midas_transforms
        del transform, img, img_original, input_batch, prediction
        return (pil2tensor(result_img), pil2tensor(depth.convert('RGB')))

    class AdjustLevels:

        def __init__(self, min_level, mid_level, max_level):
            self.min_level = min_level
            self.mid_level = mid_level
            self.max_level = max_level

        def adjust(self, im):
            im_arr = np.array(im)
            im_arr[im_arr < self.min_level] = self.min_level
            im_arr = (im_arr - self.min_level) * (255 / (self.max_level - self.min_level))
            im_arr[im_arr < 0] = 0
            im_arr[im_arr > 255] = 255
            im_arr = im_arr.astype(np.uint8)
            im = Image.fromarray(im_arr)
            im = ImageOps.autocontrast(im, cutoff=self.max_level)
            return im

    def install_midas(self):
        global MIDAS_INSTALLED
        if 'timm' not in packages():
            install_package('timm')
        MIDAS_INSTALLED = True

7.5 KiB Raw Permalink Blame History

Documentation

Input types

Required

Optional

Output types

Usage tips

Source code

7.5 KiB

Raw Permalink Blame History