From 284854a109598d29768cdfefeb7453bdc2242d9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E9=9B=AA=E5=B3=B0?= Date: Tue, 23 Apr 2024 19:02:36 +0800 Subject: [PATCH] Expand num_sort options for node InsightFaceBBOXDetect --- easyapi/DetectNode.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/easyapi/DetectNode.py b/easyapi/DetectNode.py index fdf3ac5..a779218 100644 --- a/easyapi/DetectNode.py +++ b/easyapi/DetectNode.py @@ -25,7 +25,7 @@ class InsightFaceBBOXDetect: "optional": { "num_color": ('STRING', {'default': '#FF0000'}), "num_pos": (['center', 'left-top', 'right-top', 'left-bottom', 'right-bottom', ], {}), - "num_sort": (['origin', 'reactor', ], {}), + "num_sort": (['origin', 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small'], {}), "INSIGHTFACE": ('INSIGHTFACE', {}) } } @@ -58,8 +58,18 @@ class InsightFaceBBOXDetect: img = cv2.cvtColor(np.array(tensor_to_pil(image)), cv2.COLOR_RGB2BGR) faces = model.get(img) - if num_sort == 'reactor': + if num_sort == 'reactor' or num_sort == 'left-right': faces = sorted(faces, key=lambda x: x.bbox[0]) + if num_sort == "right-left": + faces = sorted(faces, key=lambda x: x.bbox[0], reverse=True) + if num_sort == "top-bottom": + faces = sorted(faces, key=lambda x: x.bbox[1]) + if num_sort == "bottom-top": + faces = sorted(faces, key=lambda x: x.bbox[1], reverse=True) + if num_sort == "small-large": + faces = sorted(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) + if num_sort == "large-small": + faces = sorted(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True) r, g, b, a = hex_to_rgba(shape_color) n_r, n_g, n_b, n_a = hex_to_rgba(num_color)