Added sort=None to SEGS order filter (#944)

* added pick first N SEGs node

* remove take first N segs node and update segs order filter to add order=None

* add comments

* update for readibility

* format code

* remove match statement
This commit is contained in:
chutchatut
2025-03-24 00:06:06 +09:00
committed by GitHub
parent 28cd2f70b2
commit 216d7fd60c
+28 -43
View File
@@ -13,6 +13,7 @@ from . import segs_upscaler
from comfy.cli_args import args
import math
from typing import Callable, Union
try:
from comfy_extras import nodes_differential_diffusion
@@ -504,7 +505,7 @@ class SEGSOrderedFilter:
def INPUT_TYPES(s):
return {"required": {
"segs": ("SEGS", ),
"target": (["area(=w*h)", "width", "height", "x1", "y1", "x2", "y2", "confidence"],),
"target": (["area(=w*h)", "width", "height", "x1", "y1", "x2", "y2", "confidence", "none"],),
"order": ("BOOLEAN", {"default": True, "label_on": "descending", "label_off": "ascending"}),
"take_start": ("INT", {"default": 0, "min": 0, "max": sys.maxsize, "step": 1}),
"take_count": ("INT", {"default": 1, "min": 0, "max": sys.maxsize, "step": 1}),
@@ -517,51 +518,35 @@ class SEGSOrderedFilter:
CATEGORY = "ImpactPack/Util"
@staticmethod
def get_sort_key_fn(target: str) -> Union[Callable, None]:
if target == "none":
return None
def sort_key_fn(seg):
x1, y1, x2, y2 = seg.crop_region
if target == "confidence": return seg.confidence
if target == "area(=w*h)": return (x2 - x1) * (y2 - y1)
if target == "width": return x2 - x1
if target == "height": return y2 - y1
if target == "x1": return x1
if target == "y1": return y1
if target == "x2": return x2
if target == "y2": return y2
raise Exception(f"[Impact Pack] SEGSOrderedFilter - Unexpected target '{target}'")
return sort_key_fn
def doit(self, segs, target, order, take_start, take_count):
segs_with_order = []
sort_key_fn = SEGSOrderedFilter.get_sort_key_fn(target)
for seg in segs[1]:
x1 = seg.crop_region[0]
y1 = seg.crop_region[1]
x2 = seg.crop_region[2]
y2 = seg.crop_region[3]
sorted_list = list(segs[1]) # make a shallow copy, so it does not mutate the original list when sort
if sort_key_fn is not None:
sorted_list.sort(key=sort_key_fn, reverse=order)
if target == "area(=w*h)":
value = (y2 - y1) * (x2 - x1)
elif target == "width":
value = x2 - x1
elif target == "height":
value = y2 - y1
elif target == "x1":
value = x1
elif target == "x2":
value = x2
elif target == "y1":
value = y1
elif target == "y2":
value = y2
elif target == "confidence":
value = seg.confidence
else:
raise Exception(f"[Impact Pack] SEGSOrderedFilter - Unexpected target '{target}'")
segs_with_order.append((value, seg))
if order:
sorted_list = sorted(segs_with_order, key=lambda x: x[0], reverse=True)
else:
sorted_list = sorted(segs_with_order, key=lambda x: x[0], reverse=False)
result_list = []
remained_list = []
for i, item in enumerate(sorted_list):
if take_start <= i < take_start + take_count:
result_list.append(item[1])
else:
remained_list.append(item[1])
return (segs[0], result_list), (segs[0], remained_list),
take_stop = take_start + take_count
return (segs[0], sorted_list[take_start:take_stop]), \
(segs[0], sorted_list[:take_start] + sorted_list[take_stop:]),
class SEGSRangeFilter: