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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user