|
|
import impact.core as core |
|
|
from impact.config import MAX_RESOLUTION |
|
|
import impact.segs_nodes as segs_nodes |
|
|
import numpy as np |
|
|
import impact.utils as utils |
|
|
|
|
|
class SAMDetectorCombined: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"sam_model": ("SAM_MODEL", ), |
|
|
"segs": ("SEGS", ), |
|
|
"image": ("IMAGE", ), |
|
|
"detection_hint": (["center-1", "horizontal-2", "vertical-2", "rect-4", "diamond-4", "mask-area", |
|
|
"mask-points", "mask-point-bbox", "none"],), |
|
|
"dilation": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), |
|
|
"threshold": ("FLOAT", {"default": 0.93, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"bbox_expansion": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), |
|
|
"mask_hint_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"mask_hint_use_negative": (["False", "Small", "Outter"], ) |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("MASK",) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, sam_model, segs, image, detection_hint, dilation, |
|
|
threshold, bbox_expansion, mask_hint_threshold, mask_hint_use_negative): |
|
|
return (core.make_sam_mask(sam_model, segs, image, detection_hint, dilation, |
|
|
threshold, bbox_expansion, mask_hint_threshold, mask_hint_use_negative), ) |
|
|
|
|
|
|
|
|
class SAMDetectorSegmented: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"sam_model": ("SAM_MODEL", ), |
|
|
"segs": ("SEGS", ), |
|
|
"image": ("IMAGE", ), |
|
|
"detection_hint": (["center-1", "horizontal-2", "vertical-2", "rect-4", "diamond-4", "mask-area", |
|
|
"mask-points", "mask-point-bbox", "none"],), |
|
|
"dilation": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), |
|
|
"threshold": ("FLOAT", {"default": 0.93, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"bbox_expansion": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), |
|
|
"mask_hint_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"mask_hint_use_negative": (["False", "Small", "Outter"], ) |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("MASK", "MASK") |
|
|
RETURN_NAMES = ("combined_mask", "batch_masks") |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, sam_model, segs, image, detection_hint, dilation, |
|
|
threshold, bbox_expansion, mask_hint_threshold, mask_hint_use_negative): |
|
|
combined_mask, batch_masks = core.make_sam_mask_segmented(sam_model, segs, image, detection_hint, dilation, |
|
|
threshold, bbox_expansion, mask_hint_threshold, |
|
|
mask_hint_use_negative) |
|
|
return (combined_mask, batch_masks, ) |
|
|
|
|
|
|
|
|
class BboxDetectorForEach: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"bbox_detector": ("BBOX_DETECTOR", ), |
|
|
"image": ("IMAGE", ), |
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"dilation": ("INT", {"default": 10, "min": -512, "max": 512, "step": 1}), |
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 100, "step": 0.1}), |
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}), |
|
|
"labels": ("STRING", {"multiline": True, "default": "all", "placeholder": "List the types of segments to be allowed, separated by commas"}), |
|
|
}, |
|
|
"optional": {"detailer_hook": ("DETAILER_HOOK",), } |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("SEGS", ) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, bbox_detector, image, threshold, dilation, crop_factor, drop_size, labels=None, detailer_hook=None): |
|
|
segs = bbox_detector.detect(image, threshold, dilation, crop_factor, drop_size, detailer_hook) |
|
|
|
|
|
if labels is not None and labels != '': |
|
|
labels = labels.split(',') |
|
|
if len(labels) > 0: |
|
|
segs, _ = segs_nodes.SEGSLabelFilter.filter(segs, labels) |
|
|
|
|
|
return (segs, ) |
|
|
|
|
|
|
|
|
class SegmDetectorForEach: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"segm_detector": ("SEGM_DETECTOR", ), |
|
|
"image": ("IMAGE", ), |
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"dilation": ("INT", {"default": 10, "min": -512, "max": 512, "step": 1}), |
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 100, "step": 0.1}), |
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}), |
|
|
"labels": ("STRING", {"multiline": True, "default": "all", "placeholder": "List the types of segments to be allowed, separated by commas"}), |
|
|
}, |
|
|
"optional": {"detailer_hook": ("DETAILER_HOOK",), } |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("SEGS", ) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, segm_detector, image, threshold, dilation, crop_factor, drop_size, labels=None, detailer_hook=None): |
|
|
segs = segm_detector.detect(image, threshold, dilation, crop_factor, drop_size, detailer_hook) |
|
|
|
|
|
if labels is not None and labels != '': |
|
|
labels = labels.split(',') |
|
|
if len(labels) > 0: |
|
|
segs, _ = segs_nodes.SEGSLabelFilter.filter(segs, labels) |
|
|
|
|
|
return (segs, ) |
|
|
|
|
|
|
|
|
class SegmDetectorCombined: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"segm_detector": ("SEGM_DETECTOR", ), |
|
|
"image": ("IMAGE", ), |
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"dilation": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("MASK",) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, segm_detector, image, threshold, dilation): |
|
|
mask = segm_detector.detect_combined(image, threshold, dilation) |
|
|
return (mask,) |
|
|
|
|
|
|
|
|
class BboxDetectorCombined(SegmDetectorCombined): |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"bbox_detector": ("BBOX_DETECTOR", ), |
|
|
"image": ("IMAGE", ), |
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"dilation": ("INT", {"default": 4, "min": -512, "max": 512, "step": 1}), |
|
|
} |
|
|
} |
|
|
|
|
|
def doit(self, bbox_detector, image, threshold, dilation): |
|
|
mask = bbox_detector.detect_combined(image, threshold, dilation) |
|
|
return (mask,) |
|
|
|
|
|
|
|
|
class SimpleDetectorForEach: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"bbox_detector": ("BBOX_DETECTOR", ), |
|
|
"image": ("IMAGE", ), |
|
|
|
|
|
"bbox_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"bbox_dilation": ("INT", {"default": 0, "min": -255, "max": 255, "step": 1}), |
|
|
|
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 100, "step": 0.1}), |
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}), |
|
|
|
|
|
"sub_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"sub_dilation": ("INT", {"default": 0, "min": -255, "max": 255, "step": 1}), |
|
|
"sub_bbox_expansion": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), |
|
|
|
|
|
"sam_mask_hint_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
}, |
|
|
"optional": { |
|
|
"sam_model_opt": ("SAM_MODEL", ), |
|
|
"segm_detector_opt": ("SEGM_DETECTOR", ), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("SEGS",) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
@staticmethod |
|
|
def detect(bbox_detector, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt=None, segm_detector_opt=None): |
|
|
segs = bbox_detector.detect(image, bbox_threshold, bbox_dilation, crop_factor, drop_size) |
|
|
|
|
|
if sam_model_opt is not None: |
|
|
mask = core.make_sam_mask(sam_model_opt, segs, image, "center-1", sub_dilation, |
|
|
sub_threshold, sub_bbox_expansion, sam_mask_hint_threshold, False) |
|
|
segs = core.segs_bitwise_and_mask(segs, mask) |
|
|
elif segm_detector_opt is not None: |
|
|
segm_segs = segm_detector_opt.detect(image, sub_threshold, sub_dilation, crop_factor, drop_size) |
|
|
mask = core.segs_to_combined_mask(segm_segs) |
|
|
segs = core.segs_bitwise_and_mask(segs, mask) |
|
|
|
|
|
return (segs,) |
|
|
|
|
|
|
|
|
def doit(self, bbox_detector, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt=None, segm_detector_opt=None): |
|
|
|
|
|
return SimpleDetectorForEach.detect(bbox_detector, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt, segm_detector_opt) |
|
|
|
|
|
|
|
|
class SimpleDetectorForEachPipe: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"detailer_pipe": ("DETAILER_PIPE", ), |
|
|
"image": ("IMAGE", ), |
|
|
|
|
|
"bbox_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"bbox_dilation": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), |
|
|
|
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 100, "step": 0.1}), |
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}), |
|
|
|
|
|
"sub_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"sub_dilation": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), |
|
|
"sub_bbox_expansion": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), |
|
|
|
|
|
"sam_mask_hint_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
}, |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("SEGS",) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
def doit(self, detailer_pipe, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, sam_mask_hint_threshold): |
|
|
|
|
|
model, clip, vae, positive, negative, wildcard, bbox_detector, segm_detector_opt, sam_model_opt, detailer_hook, refiner_model, refiner_clip, refiner_positive, refiner_negative = detailer_pipe |
|
|
|
|
|
return SimpleDetectorForEach.detect(bbox_detector, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt, segm_detector_opt) |
|
|
|
|
|
|
|
|
class SimpleDetectorForAnimateDiff: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"bbox_detector": ("BBOX_DETECTOR", ), |
|
|
"image_frames": ("IMAGE", ), |
|
|
|
|
|
"bbox_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"bbox_dilation": ("INT", {"default": 0, "min": -255, "max": 255, "step": 1}), |
|
|
|
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 100, "step": 0.1}), |
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}), |
|
|
|
|
|
"sub_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
"sub_dilation": ("INT", {"default": 0, "min": -255, "max": 255, "step": 1}), |
|
|
"sub_bbox_expansion": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), |
|
|
|
|
|
"sam_mask_hint_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), |
|
|
}, |
|
|
"optional": { |
|
|
"sam_model_opt": ("SAM_MODEL", ), |
|
|
"segm_detector_opt": ("SEGM_DETECTOR", ), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("SEGS",) |
|
|
FUNCTION = "doit" |
|
|
|
|
|
CATEGORY = "ImpactPack/Detector" |
|
|
|
|
|
@staticmethod |
|
|
def detect(bbox_detector, image_frames, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt=None, segm_detector_opt=None): |
|
|
|
|
|
|
|
|
all_segs = [] |
|
|
for image in image_frames: |
|
|
image = image.unsqueeze(0) |
|
|
segs = bbox_detector.detect(image, bbox_threshold, bbox_dilation, crop_factor, drop_size) |
|
|
|
|
|
if sam_model_opt is not None: |
|
|
mask = core.make_sam_mask(sam_model_opt, segs, image, "center-1", sub_dilation, |
|
|
sub_threshold, sub_bbox_expansion, sam_mask_hint_threshold, False) |
|
|
segs = core.segs_bitwise_and_mask(segs, mask) |
|
|
elif segm_detector_opt is not None: |
|
|
segm_segs = segm_detector_opt.detect(image, sub_threshold, sub_dilation, crop_factor, drop_size) |
|
|
mask = core.segs_to_combined_mask(segm_segs) |
|
|
segs = core.segs_bitwise_and_mask(segs, mask) |
|
|
|
|
|
all_segs.append(segs) |
|
|
|
|
|
|
|
|
all_masks = [] |
|
|
for segs in all_segs: |
|
|
all_masks += segs_nodes.SEGSToMaskList().doit(segs)[0] |
|
|
|
|
|
result_mask = all_masks[0] |
|
|
for mask in all_masks[1:]: |
|
|
result_mask += mask |
|
|
|
|
|
result_mask = utils.to_binary_mask(result_mask, 0.1) |
|
|
|
|
|
return segs_nodes.MaskToSEGS().doit(result_mask, False, crop_factor, False, drop_size) |
|
|
|
|
|
def doit(self, bbox_detector, image_frames, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt=None, segm_detector_opt=None): |
|
|
|
|
|
return SimpleDetectorForAnimateDiff.detect(bbox_detector, image_frames, bbox_threshold, bbox_dilation, crop_factor, drop_size, |
|
|
sub_threshold, sub_dilation, sub_bbox_expansion, |
|
|
sam_mask_hint_threshold, sam_model_opt, segm_detector_opt) |
|
|
|
