lcolokandDr.Lt.Data
95cc4ed719
* Enhance mask_to_segs to Handle Batch Masks * start to build SAMDetectorSegmented * Add Segmented SAMDetector and SEGS & MASKS ForEach nodes * Bug fix: Shape tuple dimension length problem * Enhance mask_to_segs to Handle Batch Masks * pr/38 bugfix * code cleansing * pr/38 bugfix #2 apply_mask_to_each_seg crash fix 'input is 2d mask not 3d masks' * change masks:'MASK' to masks:'MASKS' * MasksToMaskList sample added robust patch for empty SEGS --------- Co-authored-by: Dr.Lt.Data <dr.lt.data@gmail.com>
145 lines
6.3 KiB
Python
145 lines
6.3 KiB
Python
import impact.core as core
|
|
from impact.config import MAX_RESOLUTION
|
|
|
|
|
|
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": 0, "max": 255, "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": 0, "max": 255, "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", "MASKS")
|
|
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": 0, "max": 255, "step": 1}),
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 10, "step": 0.1}),
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SEGS", )
|
|
FUNCTION = "doit"
|
|
|
|
CATEGORY = "ImpactPack/Detector"
|
|
|
|
def doit(self, bbox_detector, image, threshold, dilation, crop_factor, drop_size):
|
|
segs = bbox_detector.detect(image, threshold, dilation, crop_factor, drop_size)
|
|
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": 0, "max": 255, "step": 1}),
|
|
"crop_factor": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 10, "step": 0.1}),
|
|
"drop_size": ("INT", {"min": 1, "max": MAX_RESOLUTION, "step": 1, "default": 10}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SEGS", )
|
|
FUNCTION = "doit"
|
|
|
|
CATEGORY = "ImpactPack/Detector"
|
|
|
|
def doit(self, segm_detector, image, threshold, dilation, crop_factor, drop_size):
|
|
segs = segm_detector.detect(image, threshold, dilation, crop_factor, drop_size)
|
|
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": 0, "max": 255, "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": 0, "max": 255, "step": 1}),
|
|
}
|
|
}
|
|
|
|
def doit(self, bbox_detector, image, threshold, dilation):
|
|
mask = bbox_detector.detect_combined(image, threshold, dilation)
|
|
return (mask,)
|