641 lines
26 KiB
Python
641 lines
26 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import cv2
|
|
|
|
from imgutils.detect import (
|
|
detect_person, detect_faces, detect_heads, detect_halfbody,
|
|
detect_hands, detect_eyes, detect_with_booru_yolo, detect_with_nudenet
|
|
)
|
|
from imgutils.detect.censor import detect_censors
|
|
|
|
from imgutils.detect.visual import detection_visualize
|
|
|
|
|
|
BBOX = "BBOX"
|
|
class BBOX_imagutils:
|
|
"""
|
|
A class to represent bounding boxes in imgutils detection.
|
|
This is a placeholder for future use if needed.
|
|
"""
|
|
def __init__(self, x0, y0, x1, y1,img_w,img_h, label="", confidence=0.0):
|
|
self.x0 = x0
|
|
self.y0 = y0
|
|
self.x1 = x1
|
|
self.y1 = y1
|
|
self.img_w = img_w
|
|
self.img_h = img_h
|
|
self.label = label
|
|
self.confidence = confidence
|
|
|
|
def __repr__(self):
|
|
return f"BBOX({self.x0}, {self.y0}, {self.x1}, {self.y1}, {self.img_w}, {self.img_h}, '{self.label}', {self.confidence})"
|
|
|
|
def get_bbox(self):
|
|
return (self.x0, self.y0, self.x1, self.y1)
|
|
|
|
def to_json(self):
|
|
"""Convert BBOX_imagutils to JSON serializable format"""
|
|
return {
|
|
"type": "BBOX_imagutils",
|
|
"x0": self.x0,
|
|
"y0": self.y0,
|
|
"x1": self.x1,
|
|
"y1": self.y1,
|
|
"img_w": self.img_w,
|
|
"img_h": self.img_h,
|
|
"label": self.label,
|
|
"confidence": self.confidence
|
|
}
|
|
|
|
@classmethod
|
|
def from_json(cls, data):
|
|
"""Create BBOX_imagutils from JSON data"""
|
|
if isinstance(data, dict) and data.get("type") == "BBOX_imagutils":
|
|
return cls(
|
|
x0=data["x0"],
|
|
y0=data["y0"],
|
|
x1=data["x1"],
|
|
y1=data["y1"],
|
|
img_w=data.get("img_w", 0),
|
|
img_h=data.get("img_h", 0),
|
|
label=data.get("label", ""),
|
|
confidence=data.get("confidence", 0.0)
|
|
)
|
|
return data
|
|
|
|
def __dict__(self):
|
|
"""Support for dict() conversion"""
|
|
return self.to_json()
|
|
|
|
|
|
DETECTION_FUNCTIONS = {
|
|
"Person Detection": detect_person,
|
|
"Face Detection": detect_faces,
|
|
"Head Detection": detect_heads,
|
|
"Halfbody Detection": detect_halfbody,
|
|
"Hand Detection": detect_hands,
|
|
"Eye Detection": detect_eyes,
|
|
"Censor Detection": detect_censors,
|
|
# "Booru YOLO Detection": detect_with_booru_yolo,
|
|
# "Nude Detection": detect_with_nudenet,
|
|
}
|
|
|
|
MODEL_VERSIONS = {
|
|
"Person Detection": ["v0", "v1", "v1.1"],
|
|
"Face Detection": ["v0", "v1", "v1.3", "v1.4"],
|
|
"Head Detection": [],
|
|
"Halfbody Detection": ["v1"],
|
|
"Hand Detection": ["v1"],
|
|
"Eye Detection": ["v1"],
|
|
"Censor Detection": ["v1"],
|
|
}
|
|
|
|
class ImgutilsGenericDetector:
|
|
"""
|
|
A generic imgutils detector node that unifies various detection functions.
|
|
"""
|
|
CATEGORY = "imgutils/detection"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"detection_type": (list(DETECTION_FUNCTIONS.keys()),),
|
|
"conf_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"iou_threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"draw_boxes": ("BOOLEAN", {"default": True}),
|
|
# Model specific parameters, always present for simplicity in this generic node
|
|
"level": ("STRING", {"default": "s", "options": ["n", "s"]}),
|
|
"version": ("STRING", {"default": "", "options": ["v0", "v1", "v1.1"]}),
|
|
# "model_name": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK",BBOX, )
|
|
RETURN_NAMES = ("image_with_boxes", "detection_mask","detection_bbox",)
|
|
FUNCTION = "detect"
|
|
|
|
def detect(self, image, detection_type, conf_threshold,
|
|
iou_threshold, draw_boxes, level, version, model_name=None):
|
|
# Convert ComfyUI image tensor to PIL Image
|
|
# image is in format [B, H, W, C] with values 0-1
|
|
img_tensor = image[0] # Get first image from batch
|
|
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8) # Convert to 0-255 range
|
|
img_pil = Image.fromarray(img_np)
|
|
|
|
# For visualization and mask generation, we still need numpy array
|
|
img_np_rgb = np.array(img_pil.convert("RGB"))
|
|
|
|
selected_detect_func = DETECTION_FUNCTIONS.get(detection_type)
|
|
detection_results_formatted = []
|
|
output_image_tensor = image # Default to original image tensor
|
|
output_mask_np = np.zeros((img_np_rgb.shape[0], img_np_rgb.shape[1]), dtype=np.float32)
|
|
output_string = ""
|
|
|
|
if not selected_detect_func:
|
|
output_string = f"Error: Detection type '{detection_type}' not supported."
|
|
output_mask_tensor = torch.from_numpy(output_mask_np).unsqueeze(0) # Add batch dimension
|
|
return (output_image_tensor, output_string, output_mask_tensor)
|
|
|
|
func_params = {
|
|
"image": img_pil, # Use PIL Image directly for imgutils functions
|
|
"conf_threshold": conf_threshold,
|
|
"iou_threshold": iou_threshold,
|
|
}
|
|
|
|
func_params["level"] = level
|
|
|
|
if version in MODEL_VERSIONS[detection_type]:
|
|
func_params["version"] = version
|
|
if model_name is None or model_name == "":
|
|
pass
|
|
else:
|
|
func_params["model_name"] = model_name
|
|
|
|
# if 1 + 1 == 2:
|
|
try:
|
|
# print(f"Running {detection_type} with parameters:")
|
|
# print(f" Original image type: {type(img_pil)}")
|
|
# print(f" Original image mode: {img_pil.mode}")
|
|
# print(f" Original image size: {img_pil.size}")
|
|
|
|
# Ensure image is in RGB mode
|
|
if img_pil.mode != 'RGB':
|
|
img_pil = img_pil.convert('RGB')
|
|
|
|
print(f"Using detection function: {selected_detect_func.__name__}")
|
|
|
|
# # Try different calling patterns based on the function
|
|
# if detection_type in ["Face Detection", "Person Detection"]:
|
|
# # These functions might have different parameter requirements
|
|
# if detection_type == "Face Detection":
|
|
# # Try with minimal parameters first
|
|
# raw_detections = selected_detect_func(
|
|
# image=img_pil,
|
|
# conf_threshold=conf_threshold
|
|
# )
|
|
# else: # Person Detection
|
|
# raw_detections = selected_detect_func(
|
|
# image=img_pil,
|
|
# level=level,
|
|
# version=version if version in MODEL_VERSIONS[detection_type] else "v1.1",
|
|
# conf_threshold=conf_threshold
|
|
# )
|
|
# else:
|
|
# # For other detection types, use the standard parameters
|
|
# raw_detections = selected_detect_func(
|
|
# image=img_pil,
|
|
# level=level,
|
|
# conf_threshold=conf_threshold
|
|
# )
|
|
|
|
if "version" in func_params:
|
|
raw_detections = selected_detect_func(
|
|
image=img_pil,
|
|
level=level,
|
|
version=func_params["version"],
|
|
conf_threshold=conf_threshold,
|
|
iou_threshold=iou_threshold
|
|
)
|
|
else:
|
|
raw_detections = selected_detect_func(
|
|
image=img_pil,
|
|
level=level,
|
|
conf_threshold=conf_threshold,
|
|
iou_threshold=iou_threshold
|
|
)
|
|
img_w, img_h = img_pil.size
|
|
|
|
for item in raw_detections:
|
|
if isinstance(item, tuple) and len(item) == 3:
|
|
box, label, confidence = item
|
|
if isinstance(box, tuple) and len(box) == 4:
|
|
detection_results_formatted.append({
|
|
'box': [int(b) for b in box],
|
|
'label': str(label),
|
|
'confidence': float(confidence)
|
|
})
|
|
bbox_list = []
|
|
|
|
if detection_results_formatted:
|
|
# output_string = f"{detection_type} Results:\n"
|
|
# output_string += "\n".join([
|
|
# f"- {res['label']}: Confidence={res['confidence']:.2f}, Box={res['box']}"
|
|
# for res in detection_results_formatted if res['confidence'] >= conf_threshold # Apply conf_threshold for string output too
|
|
# ])
|
|
# print(f"Detection results: {output_string}")
|
|
for res in detection_results_formatted:
|
|
if res['confidence'] >= conf_threshold:
|
|
bbox = BBOX_imagutils(
|
|
x0=res['box'][0],
|
|
y0=res['box'][1],
|
|
x1=res['box'][2],
|
|
y1=res['box'][3],
|
|
img_w=img_w,
|
|
img_h=img_h,
|
|
label=res['label'],
|
|
confidence=res['confidence']
|
|
)
|
|
bbox_list.append(bbox)
|
|
else:
|
|
print(f"No {detection_type} detected with confidence >= {conf_threshold}.")
|
|
# output_string = f"No {detection_type} detected with confidence >= {conf_threshold}."
|
|
|
|
|
|
if draw_boxes and detection_results_formatted:
|
|
vis_input_for_detection_visualize = []
|
|
for res in detection_results_formatted:
|
|
if res['confidence'] >= conf_threshold:
|
|
vis_input_for_detection_visualize.append(
|
|
(tuple(res['box']), res['label'], res['confidence'])
|
|
)
|
|
|
|
if vis_input_for_detection_visualize: # Only visualize if there are results after filtering
|
|
output_image_pil = detection_visualize(
|
|
image=img_pil,
|
|
detection=vis_input_for_detection_visualize,
|
|
fontsize=12,
|
|
no_label=False
|
|
)
|
|
# Convert PIL back to ComfyUI tensor format
|
|
output_image_np = np.array(output_image_pil.convert("RGB"))
|
|
output_image_tensor = torch.from_numpy(output_image_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
else:
|
|
output_image_tensor = image # Use original image tensor
|
|
|
|
else: # If not drawing or no results, use original image tensor
|
|
output_image_tensor = image # Generate mask
|
|
for res in detection_results_formatted:
|
|
if res['confidence'] >= conf_threshold:
|
|
box = res['box']
|
|
cv2.rectangle(output_mask_np, (box[0], box[1]), (box[2], box[3]), 1.0, -1)
|
|
|
|
# Convert mask to ComfyUI tensor format [B, H, W]
|
|
output_mask_tensor = torch.from_numpy(output_mask_np).unsqueeze(0)
|
|
|
|
return (
|
|
output_image_tensor,
|
|
output_mask_tensor,
|
|
bbox_list if bbox_list else [],
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Error during {detection_type} detection: {e}")
|
|
# Error handling - return original image, error message, and empty mask
|
|
# output_string = f"Error during '{detection_type}' detection: {e}"
|
|
empty_mask = torch.zeros((1, img_np_rgb.shape[0], img_np_rgb.shape[1]), dtype=torch.float32)
|
|
empty_bbox_list = []
|
|
return (image, empty_mask, empty_bbox_list)
|
|
|
|
|
|
class MaskToBBoxNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Threshold to binarize the mask before finding bbox"}),
|
|
"combine": ("BOOLEAN", {"default": False, "tooltip": "If True, combine all mask regions into one bbox. If False, output separate bboxes for each connected component."}),
|
|
"nms_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "IoU threshold for Non-Maximum Suppression. Higher values merge more overlapping boxes."})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOX",)
|
|
FUNCTION = "mask_to_bbox"
|
|
CATEGORY = "imgutils/bbox"
|
|
DESCRIPTION = "Calculates the bounding box(es) of a mask. Can combine all regions or separate them with NMS."
|
|
|
|
def calculate_iou(self, box1, box2):
|
|
"""
|
|
Calculate Intersection over Union (IoU) of two bounding boxes.
|
|
|
|
Args:
|
|
box1, box2: [x_min, y_min, x_max, y_max]
|
|
|
|
Returns:
|
|
IoU value between 0 and 1
|
|
"""
|
|
x1_min, y1_min, x1_max, y1_max = box1
|
|
x2_min, y2_min, x2_max, y2_max = box2
|
|
|
|
# Calculate intersection area
|
|
inter_x_min = max(x1_min, x2_min)
|
|
inter_y_min = max(y1_min, y2_min)
|
|
inter_x_max = min(x1_max, x2_max)
|
|
inter_y_max = min(y1_max, y2_max)
|
|
|
|
if inter_x_max <= inter_x_min or inter_y_max <= inter_y_min:
|
|
return 0.0
|
|
|
|
inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min)
|
|
|
|
# Calculate union area
|
|
box1_area = (x1_max - x1_min) * (y1_max - y1_min)
|
|
box2_area = (x2_max - x2_min) * (y2_max - y2_min)
|
|
union_area = box1_area + box2_area - inter_area
|
|
|
|
if union_area == 0:
|
|
return 0.0
|
|
|
|
return inter_area / union_area
|
|
|
|
def merge_boxes(self, box1, box2):
|
|
"""
|
|
Merge two bounding boxes by taking the outer bounds.
|
|
|
|
Args:
|
|
box1, box2: [x_min, y_min, x_max, y_max]
|
|
|
|
Returns:
|
|
Merged bounding box [x_min, y_min, x_max, y_max]
|
|
"""
|
|
x1_min, y1_min, x1_max, y1_max = box1
|
|
x2_min, y2_min, x2_max, y2_max = box2
|
|
|
|
return [
|
|
min(x1_min, x2_min),
|
|
min(y1_min, y2_min),
|
|
max(x1_max, x2_max),
|
|
max(y1_max, y2_max)
|
|
]
|
|
|
|
def apply_nms(self, bbox_list, nms_threshold):
|
|
"""
|
|
Apply Non-Maximum Suppression to merge overlapping bounding boxes.
|
|
|
|
Args:
|
|
bbox_list: List of bounding boxes [[x_min, y_min, x_max, y_max], ...]
|
|
nms_threshold: IoU threshold for merging
|
|
|
|
Returns:
|
|
List of merged bounding boxes
|
|
"""
|
|
if len(bbox_list) <= 1:
|
|
return bbox_list
|
|
|
|
merged_boxes = []
|
|
used = [False] * len(bbox_list)
|
|
img_w = bbox_list[0].img_w if isinstance(bbox_list[0], BBOX_imagutils) else 1
|
|
img_h = bbox_list[0].img_h if isinstance(bbox_list[0], BBOX_imagutils) else 1
|
|
|
|
for i in range(len(bbox_list)):
|
|
if used[i]:
|
|
continue
|
|
|
|
current_box = bbox_list[i]
|
|
boxes_to_merge = [current_box]
|
|
used[i] = True
|
|
|
|
# Find all boxes that overlap with current box above threshold
|
|
# for j in range(i + 1, len(bbox_list)):
|
|
# if used[j]:
|
|
# continue
|
|
|
|
# iou = self.calculate_iou(current_box, bbox_list[j])
|
|
# if iou >= nms_threshold:
|
|
# boxes_to_merge.append(bbox_list[j])
|
|
# used[j] = True
|
|
|
|
# # Merge all overlapping boxes
|
|
# if len(boxes_to_merge) == 1:
|
|
# merged_boxes.append(current_box)
|
|
# else:
|
|
# merged_box = boxes_to_merge[0]
|
|
# for box in boxes_to_merge[1:]:
|
|
# merged_box = self.merge_boxes(merged_box, box)
|
|
# merged_boxes.append(merged_box)
|
|
|
|
for j in range(i + 1, len(bbox_list)):
|
|
if used[j]:
|
|
continue
|
|
|
|
# 取出坐标用于IoU计算
|
|
iou = self.calculate_iou(current_box.get_bbox(), bbox_list[j].get_bbox())
|
|
if iou >= nms_threshold:
|
|
boxes_to_merge.append(bbox_list[j])
|
|
used[j] = True
|
|
|
|
# Merge all overlapping boxes
|
|
if len(boxes_to_merge) == 1:
|
|
merged_boxes.append(current_box)
|
|
else:
|
|
# 合并所有box，取最大外接框，label和confidence可自定义合并策略
|
|
merged_box_coords = boxes_to_merge[0].get_bbox()
|
|
label = boxes_to_merge[0].label
|
|
confidence = max(box.confidence for box in boxes_to_merge)
|
|
for box in boxes_to_merge[1:]:
|
|
merged_box_coords = self.merge_boxes(merged_box_coords, box.get_bbox())
|
|
merged_box = BBOX_imagutils(
|
|
x0=merged_box_coords[0],
|
|
y0=merged_box_coords[1],
|
|
x1=merged_box_coords[2],
|
|
y1=merged_box_coords[3],
|
|
img_w=img_w,
|
|
img_h=img_h,
|
|
label=label,
|
|
confidence=confidence
|
|
)
|
|
merged_boxes.append(merged_box)
|
|
|
|
return merged_boxes
|
|
|
|
def mask_to_bbox(self, mask, threshold=0.5, combine=True, nms_threshold=0.5):
|
|
if mask.shape[0] > 1:
|
|
print(f"Warning: MaskToBBoxNode received a batch of masks ({mask.shape[0]}). Calculating bbox for the first mask only.")
|
|
|
|
# Convert to boolean numpy array (H, W)
|
|
mask_np = mask[0].squeeze(0).cpu().numpy() > threshold
|
|
mask_w, mask_h = mask_np.shape
|
|
|
|
if combine:
|
|
# Original behavior - combine all mask regions into one bbox
|
|
rows = np.any(mask_np, axis=1)
|
|
cols = np.any(mask_np, axis=0)
|
|
|
|
if not np.any(rows) or not np.any(cols):
|
|
print("Warning: Mask is empty or fully transparent, cannot calculate bbox.")
|
|
bbox = [0, 0, 0, 0]
|
|
else:
|
|
y_min, y_max = np.where(rows)[0][[0, -1]]
|
|
x_min, x_max = np.where(cols)[0][[0, -1]]
|
|
bbox = [int(x_min), int(y_min), int(x_max), int(y_max)]
|
|
|
|
print(f"Calculated combined BBox: {bbox}")
|
|
return (bbox,)
|
|
|
|
else:
|
|
from scipy import ndimage
|
|
|
|
# Label connected components
|
|
labeled_mask, num_components = ndimage.label(mask_np)
|
|
|
|
if num_components == 0:
|
|
print("Warning: Mask is empty or fully transparent, cannot calculate bbox.")
|
|
bbox_list = []
|
|
else:
|
|
bbox_list = []
|
|
|
|
for component_id in range(1, num_components + 1):
|
|
# Get mask for this component
|
|
component_mask = (labeled_mask == component_id)
|
|
|
|
# Find bounding box for this component
|
|
rows = np.any(component_mask, axis=1)
|
|
cols = np.any(component_mask, axis=0)
|
|
|
|
if np.any(rows) and np.any(cols):
|
|
y_min, y_max = np.where(rows)[0][[0, -1]]
|
|
x_min, x_max = np.where(cols)[0][[0, -1]]
|
|
# bbox = [int(x_min), int(y_min), int(x_max), int(y_max)]
|
|
# bbox = [int(x_min), int(y_min), int(x_max), int(y_max),""]
|
|
bbox = BBOX_imagutils(
|
|
x0=int(x_min),
|
|
y0=int(y_min),
|
|
x1=int(x_max),
|
|
y1=int(y_max),
|
|
img_w=mask_w,
|
|
img_h=mask_h,
|
|
label="", # No label for generic bbox
|
|
confidence=1.0 # Confidence is 1.0 for mask-based bboxes
|
|
)
|
|
bbox_list.append(bbox)
|
|
|
|
# Apply NMS to merge overlapping boxes
|
|
if bbox_list and nms_threshold < 1.0:
|
|
bbox_list = self.apply_nms(bbox_list, nms_threshold)
|
|
print(f"After NMS (threshold={nms_threshold}): {len(bbox_list)} BBoxes remain")
|
|
|
|
print(f"Final BBoxes: {bbox_list}")
|
|
|
|
return (bbox_list,)
|
|
|
|
class BBoxToMaskNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
# "image_shape": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1, "tooltip": "Height and width of the output mask"}),
|
|
"merge_options": ("BOOLEAN", {"default": False, "tooltip": "If True, merge all bounding boxes into one mask. If False, create separate masks for each bbox."}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "bbox_to_mask"
|
|
|
|
CATEGORY = "imgutils/bbox"
|
|
DESCRIPTION = "Converts bounding boxes to a binary mask."
|
|
|
|
def bbox_to_mask(self, bbox, merge_options=False, image_shape=512):
|
|
|
|
mask_list = []
|
|
bbox_input_list = []
|
|
if bbox is not None:
|
|
if isinstance(bbox, list) and len(bbox) > 0:
|
|
if isinstance(bbox[0], BBOX_imagutils):
|
|
bbox_input_list = [b.get_bbox() for b in bbox]
|
|
elif isinstance(bbox, BBOX_imagutils):
|
|
bbox_input_list = [bbox.get_bbox()]
|
|
|
|
if len(bbox_input_list) == 0:
|
|
print("Warning: No bounding boxes provided, returning empty mask.")
|
|
return (torch.zeros((1, image_shape, image_shape), dtype=torch.float32),)
|
|
|
|
mask_w, mask_h = image_shape, image_shape
|
|
|
|
if merge_options:
|
|
# make sure all boxes are in same image shape
|
|
img_w, img_h = bbox_input_list[0].img_w, bbox_input_list[0].img_h
|
|
for box in bbox_input_list:
|
|
if box.img_w != img_w or box.img_h != img_h:
|
|
|
|
print(f"Warning: Bounding boxes have different image shapes ({box.img_w}x{box.img_h}), using first box shape ({img_w}x{img_h}) for mask.")
|
|
break
|
|
mask_w, mask_h = img_w, img_h
|
|
|
|
for box in bbox_input_list:
|
|
x0, y0, x1, y1 = box.get_bbox()
|
|
# Ensure coordinates are within bounds
|
|
x0 = max(0, min(x0, mask_w - 1))
|
|
y0 = max(0, min(y0, mask_h - 1))
|
|
x1 = max(0, min(x1, mask_w - 1))
|
|
y1 = max(0, min(y1, mask_h - 1))
|
|
|
|
mask = torch.zeros((1, mask_h, mask_w), dtype=torch.float32)
|
|
mask[0, y0:y1, x0:x1] = 1.
|
|
mask_list.append(mask)
|
|
|
|
if merge_options:
|
|
# Merge all masks into one
|
|
merged_mask = torch.zeros((1, mask_h, mask_w), dtype=torch.float32)
|
|
for mask in mask_list:
|
|
merged_mask += mask
|
|
merged_mask = (merged_mask > 0).float()
|
|
|
|
return (merged_mask,)
|
|
else:
|
|
return (mask_list,)
|
|
|
|
class BBoxFilter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"bboxes": (BBOX,),
|
|
"labels": ("STRING", {"default": "", "tooltip": "Comma-separated list of labels to filter by. Leave empty to keep all labels."}),
|
|
"include_labels": ("BOOLEAN", {"default": True, "tooltip": "If True, only keep bounding boxes with specified labels. If False, exclude these labels."}),
|
|
"min_area": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000000.0, "step": 1.0, "tooltip": "Minimum area of bounding boxes to keep"}),
|
|
"min_confidence": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Minimum confidence of bounding boxes to keep"}),
|
|
}
|
|
}
|
|
RETURN_TYPES = (BBOX,)
|
|
RETURN_NAMES = ("filtered_bboxes",)
|
|
FUNCTION = "filter_bboxes"
|
|
|
|
CATEGORY = "imgutils/bbox"
|
|
DESCRIPTION = "Filters bounding boxes based on area, confidence, and labels."
|
|
def filter_bboxes(self, bboxes,include_labels, min_area=0.0, min_confidence=0.0, labels=""):
|
|
|
|
filtered_bboxes = []
|
|
# Labels 应该是一个逗号分隔的字符串，如果 labels 为空或者仅有空格，则表示不过滤标签
|
|
if labels:
|
|
labels = [label.strip() for label in labels.split(",") if label.strip()]
|
|
else:
|
|
labels = None
|
|
|
|
for bbox in bboxes:
|
|
if isinstance(bbox, BBOX_imagutils):
|
|
area = (bbox.x1 - bbox.x0) * (bbox.y1 - bbox.y0)
|
|
if area < min_area or bbox.confidence < min_confidence:
|
|
continue
|
|
if include_labels:
|
|
if labels and bbox.label not in labels:
|
|
continue
|
|
else: # Exclude labels
|
|
if labels and bbox.label in labels:
|
|
continue
|
|
|
|
filtered_bboxes.append(bbox)
|
|
else:
|
|
print(f"Warning: Skipping non-BBOX_imagutils object: {bbox}")
|
|
|
|
return (filtered_bboxes,)
|
|
|
|
|
|
# ComfyUI Node Registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImgutilsGenericDetector": ImgutilsGenericDetector,
|
|
"MaskToBBoxNode": MaskToBBoxNode,
|
|
"BBoxToMaskNode": BBoxToMaskNode,
|
|
"BBoxFilter": BBoxFilter,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImgutilsGenericDetector": "Imgutils Generic Detector",
|
|
"MaskToBBoxNode": "Mask to BBox",
|
|
"BBoxToMaskNode": "BBox to Mask",
|
|
"BBoxFilter": "BBox Filter",
|
|
}
|
|
|
|
|