From 0f3303f439a88f585561b9fbdb449e6c138ecc45 Mon Sep 17 00:00:00 2001
From: SayanoAI <28544690+SayanoAI@users.noreply.github.com>
Date: Thu, 8 Aug 2024 14:56:21 -0400
Subject: [PATCH] sort images node

---
 __init__.py           |  6 ++++--
 custom_nodes/utils.py | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/__init__.py b/__init__.py
index d96c169..da1a114 100644
--- a/__init__.py
+++ b/__init__.py
@@ -3,7 +3,7 @@ from .custom_nodes.uvr import UVR5Node
 from .custom_nodes.rvc import RVCNode
 from .custom_nodes.loaders import DownloadAudio, LoadAudio, LoadWhisperModelNode, LoadRVCModelNode, LoadHubertModel, LoadPitchExtractionParams
 from .custom_nodes.output import PreviewAudio
-from .custom_nodes.utils import Any2ListNode, AudioBatchValueNode, List2AnyNode, MergeImageBatches, MergeLatentBatches, ImageRepeatInterleavedNode, LatentRepeatInterleavedNode, MergeAudioNode, SimpleMathNode, SliceNode, ZipImagesNode
+from .custom_nodes.utils import Any2ListNode, AudioBatchValueNode, List2AnyNode, MergeImageBatches, MergeLatentBatches, ImageRepeatInterleavedNode, LatentRepeatInterleavedNode, MergeAudioNode, SimpleMathNode, SliceNode, SortImagesNode, ZipImagesNode
 
 # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
 WEB_DIRECTORY = "./web"
@@ -32,7 +32,8 @@ NODE_CLASS_MAPPINGS = {
     "SliceNode": SliceNode,
     "ZipNode": ZipImagesNode,
     "Any2ListNode": Any2ListNode,
-    "List2AnyNode": List2AnyNode
+    "List2AnyNode": List2AnyNode,
+    "SortImagesNode": SortImagesNode
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -59,4 +60,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ZipNode": "🌺Zip Images",
     "Any2ListNode": "🌺Any to List",
     "List2AnyNode": "🌺List to Any",
+    "SortImagesNode": "🌺Sort Images",
 }
\ No newline at end of file
diff --git a/custom_nodes/utils.py b/custom_nodes/utils.py
index 4595756..b108664 100644
--- a/custom_nodes/utils.py
+++ b/custom_nodes/utils.py
@@ -548,4 +548,39 @@ class List2AnyNode:
     CATEGORY = CATEGORY
 
     def to(self, any):
-        return (any,)
\ No newline at end of file
+        return (any,)
+    
+class SortImagesNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            },
+            "optional": {
+                "indices": ("INT", {"forceInput": True}),
+                "reverse": ("BOOLEAN", {"default": False}),
+                "sort_by": (["sum","mean","median","min","max"],{"default": "sum"})
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "INT")
+    RETURN_NAMES = ("images", "indices")
+    FUNCTION = "execute"
+    CATEGORY = CATEGORY
+
+    def execute(self, images, indices=None, reverse=False, sort_by="sum"):
+        if sort_by=="mean": func=np.mean
+        elif sort_by=="median": func=np.median
+        elif sort_by=="min": func=np.amin
+        elif sort_by=="max": func=np.amax
+        else: func=np.sum
+        values = list(map(lambda x: func(x.numpy()),images))
+        if indices is None:
+            indices = np.argsort(values)
+        if reverse: indices=indices[::-1]
+        indices = list(indices)
+        return (images[indices],indices)
\ No newline at end of file