110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
import impact_core as core
|
|
|
|
|
|
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 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}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SEGS", )
|
|
FUNCTION = "doit"
|
|
|
|
CATEGORY = "ImpactPack/Detector"
|
|
|
|
def doit(self, bbox_detector, image, threshold, dilation, crop_factor):
|
|
segs = bbox_detector.detect(image, threshold, dilation, crop_factor)
|
|
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}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SEGS", )
|
|
FUNCTION = "doit"
|
|
|
|
CATEGORY = "ImpactPack/Detector"
|
|
|
|
def doit(self, segm_detector, image, threshold, dilation, crop_factor):
|
|
segs = segm_detector.detect(image, threshold, dilation, crop_factor)
|
|
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,)
|