|
import torch |
|
from PIL import Image |
|
from rembg import remove, new_session |
|
|
|
import folder_paths |
|
|
|
from ..session.CustomSession import CustomAbstractSession |
|
from ..session.CustomSession import CustomSessionContainer |
|
from ..session.ModnetPhotographicSession import ModnetPhotographicSession |
|
from ..session.ModnetWebcamSession import ModnetWebcamSession |
|
|
|
|
|
class ImageSegmentation: |
|
def __init__(self): |
|
pass |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
"required": { |
|
"images": ("IMAGE",), |
|
"model": ([ |
|
"u2net", |
|
"u2netp", |
|
"u2net_human_seg", |
|
"u2net_cloth_seg", |
|
"silueta", |
|
"isnet-general-use", |
|
"isnetis", |
|
"modnet-p", |
|
"modnet-w" |
|
],), |
|
"alpha_matting": (["true", "false"],), |
|
"alpha_matting_foreground_threshold": ("INT", { |
|
"default": 240, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_background_threshold": ("INT", { |
|
"default": 20, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_erode_size": ("INT", { |
|
"default": 10, |
|
"step": 1 |
|
}), |
|
"post_process_mask": (["false", "true"],), |
|
}, |
|
} |
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
FUNCTION = "node" |
|
CATEGORY = "image/segmentation" |
|
|
|
def node( |
|
self, |
|
images, |
|
model, |
|
alpha_matting, |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, |
|
post_process_mask, |
|
session=None |
|
): |
|
if session is None: |
|
if model == "isnetis": |
|
session = new_session("isnet-anime") |
|
elif model == "modnet-p": |
|
session = ModnetPhotographicSession(model) |
|
elif model == "modnet-w": |
|
session = ModnetWebcamSession(model) |
|
else: |
|
session = new_session(model) |
|
|
|
def verst(image): |
|
img: Image = image.tensor_to_image() |
|
|
|
return remove( |
|
img, alpha_matting == "true", |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, session, |
|
False, post_process_mask == "true" |
|
).image_to_tensor() |
|
|
|
return (torch.stack([ |
|
verst(images[i]) for i in range(len(images)) |
|
]),) |
|
|
|
|
|
class ImageSegmentationCustom: |
|
def __init__(self): |
|
pass |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
"required": { |
|
"images": ("IMAGE",), |
|
"model": (folder_paths.get_filename_list("onnx"),), |
|
"alpha_matting": (["true", "false"],), |
|
"alpha_matting_foreground_threshold": ("INT", { |
|
"default": 240, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_background_threshold": ("INT", { |
|
"default": 20, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_erode_size": ("INT", { |
|
"default": 10, |
|
"step": 1 |
|
}), |
|
"post_process_mask": (["false", "true"],), |
|
"mean": ("FLOAT", { |
|
"default": 0.485, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"std": ("FLOAT", { |
|
"default": 1.0, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"size": ("INT", { |
|
"default": 1024, |
|
"step": 8 |
|
}), |
|
}, |
|
} |
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
FUNCTION = "node" |
|
CATEGORY = "image/segmentation" |
|
|
|
def node( |
|
self, |
|
images, |
|
model, |
|
alpha_matting, |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, |
|
post_process_mask, |
|
mean, |
|
std, |
|
size |
|
): |
|
container = CustomSessionContainer(mean, mean, mean, std, std, std, size, size) |
|
|
|
class CustomSession(CustomAbstractSession): |
|
def __init__(self): |
|
super().__init__(model) |
|
|
|
@classmethod |
|
def name(cls, *args, **kwargs): |
|
return model |
|
|
|
session = CustomSession().from_container(container) |
|
|
|
return ImageSegmentation().node( |
|
images, |
|
model, |
|
alpha_matting, |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, |
|
post_process_mask, |
|
session |
|
) |
|
|
|
|
|
class ImageSegmentationCustomAdvanced: |
|
def __init__(self): |
|
pass |
|
|
|
@classmethod |
|
def INPUT_TYPES(cls): |
|
return { |
|
"required": { |
|
"images": ("IMAGE",), |
|
"model": (folder_paths.get_filename_list("onnx"),), |
|
"alpha_matting": (["true", "false"],), |
|
"alpha_matting_foreground_threshold": ("INT", { |
|
"default": 240, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_background_threshold": ("INT", { |
|
"default": 20, |
|
"max": 250, |
|
"step": 5 |
|
}), |
|
"alpha_matting_erode_size": ("INT", { |
|
"default": 10, |
|
"step": 1 |
|
}), |
|
"post_process_mask": (["false", "true"],), |
|
"mean_r": ("FLOAT", { |
|
"default": 0.485, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"mean_g": ("FLOAT", { |
|
"default": 0.456, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"mean_b": ("FLOAT", { |
|
"default": 0.406, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"std_r": ("FLOAT", { |
|
"default": 1.0, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"std_g": ("FLOAT", { |
|
"default": 1.0, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"std_b": ("FLOAT", { |
|
"default": 1.0, |
|
"max": 1.0, |
|
"step": 0.01 |
|
}), |
|
"width": ("INT", { |
|
"default": 1024, |
|
"step": 8 |
|
}), |
|
"height": ("INT", { |
|
"default": 1024, |
|
"step": 8 |
|
}), |
|
}, |
|
} |
|
|
|
RETURN_TYPES = ("IMAGE",) |
|
FUNCTION = "node" |
|
CATEGORY = "image/segmentation" |
|
|
|
def node( |
|
self, |
|
images, |
|
model, |
|
alpha_matting, |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, |
|
post_process_mask, |
|
mean_x, |
|
mean_y, |
|
mean_z, |
|
std_x, |
|
std_y, |
|
std_z, |
|
width, |
|
height |
|
): |
|
container = CustomSessionContainer(mean_x, mean_y, mean_z, std_x, std_y, std_z, width, height) |
|
|
|
class CustomSession(CustomAbstractSession): |
|
def __init__(self): |
|
super().__init__(model) |
|
|
|
@classmethod |
|
def name(cls, *args, **kwargs): |
|
return model |
|
|
|
session = CustomSession().from_container(container) |
|
|
|
return ImageSegmentation().node( |
|
images, |
|
model, |
|
alpha_matting, |
|
alpha_matting_foreground_threshold, |
|
alpha_matting_background_threshold, |
|
alpha_matting_erode_size, |
|
post_process_mask, |
|
session |
|
) |
|
|
|
|
|
NODE_CLASS_MAPPINGS = { |
|
"ImageSegmentation": ImageSegmentation, |
|
"ImageSegmentationCustom": ImageSegmentationCustom, |
|
"ImageSegmentationCustomAdvanced": ImageSegmentationCustomAdvanced, |
|
} |
|
|
