Advanced Result Browser

README.md

@@ -20,15 +20,19 @@ This repository provides several custom nodes for ComfyUI that enable efficient
 - new_db_name: Choose a filename with the .json extension. The file will automatically be stored in your ComfyUI "/models/embDBs" folder.
 - If you have Crystools installed, you can use the "Show any value to console/display" node to display a list of files that could not be recorded in the database (any other node capable of displaying text output works as well).
 
-Additionally you can chain up as many databases has you want by using the optional "img_db" socket.
-
 ## After creating a database, you can search for similar images using a workflow like the one shown above:
 ![Using a ClipVision-Database](examples/UsingClipVision-Database.png)
 - Load an image into the Get Image Embeddings node, connected to a "Load CLIP Vision" node and select a Clip Vision model. → This node generates the embedding data of your input image.
 - Pass this embedding data to the "Image Searcher" node, which calculates image similarities using a cosine similarity algorithm.
-- Browse the search results with the "Result Browser" node:
+- "Result Browser": Browse the search results
   - image_index: Index of the image, depending on match direction.
   - match: "first" / "last" - "first" → image_index 0 outputs the most similar image. "last" → image_index 0 outputs the least similar image.
+- "Advanced Result Browser": Similar to "Result Browser", but extended to output an image batch, all images scaled to resolution of "batch_frame_image" or resolution of the first image of the batch.
+  - offset_index: see image_index of "Result Browser"
+  - image_count: the amount of images to output as a image batch
+  - match: see match of "Result Browser"
+
+Additionally you can chain up as many databases has you want by using the optional "img_db" socket.
 
 ## You can directly compare the similarity of two images by connecting the outputs of two "Get Image Embeddings" nodes to a "Compare Embeddings" node.
 ![Simple image comparison](examples/SimpleImageComparison.png)

__init__.py

@@ -22,9 +22,9 @@ NODE_CLASS_MAPPINGS = {
     "CalcEmbeds": CalcEmbeds,
     "CompareEmbeds": CompareEmbeds,
     "ImageSearcher": ImageSearcher,
-    "RBrowser": ResultBrowser,
-    "AdvRBrowser": ResultBrowserAdvanced,
-    "RCombiner": ResultCombiner,
+    "ResultBrowser": ResultBrowser,
+    "ResultBrowserAdvanced": ResultBrowserAdvanced,
+    "ResultCombiner": ResultCombiner,
     "LoadDB": LoadDB,
     "GenerateDB": GenerateDB,
 }
@@ -36,9 +36,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CalcEmbeds": "Calculate Embeddings (experimental)",
     "CompareEmbeds": "Compare Embeddings",
     "ImageSearcher": "Image Searcher",
-    "RBrowser": "Result Browser",
-    "AdvRBrowser": "Advanced Result Browser",
-    "RCombiner": "Combine Results",
+    "ResultBrowser": "Result Browser",
+    "ResultBrowserAdvanced": "Advanced Result Browser",
+    "ResultCombiner": "Combine Results",
     "DBLoader": "Embeddings DB Loader",
     "GenerateDB": "Generate Image Database"
 }

clipvision_db.py

@@ -1,5 +1,4 @@
 import os
-import json
 from pathlib import Path
 import folder_paths
 import orjson
@@ -7,6 +6,10 @@ import orjson
 from .utils import (generate_clip_features_json)
 
 class LoadDB:
+    """
+    Loads an image database from a JSON file.
+    """
+
     MY_LOADED_DB: list = None
     LOADED_DB: list = None
     LOADED_DB_Name = ""
@@ -36,18 +39,33 @@ class LoadDB:
     CATEGORY = "ClipVisionTools"
 
     def load_DB(self, db_name, path_to_images_folder, img_db=None):
+        """
+        Loads the image database from the specified JSON file.
+
+        Args:
+            db_name: The name of the database file.
+            path_to_images_folder: The path to the folder containing the images.
+            img_db: An optional existing LoadDB object to append to.
+
+        Returns:
+            A LoadDB object.
+        """        
         db_path = folder_paths.get_full_path_or_raise("EmbDBs", db_name)
         if self.LOADED_DB_Name != db_path or path_to_images_folder != self.images_folder:
-            with open(db_path, "rb") as f:
-                tmp = orjson.loads(f.read())
-
-            self.MY_LOADED_DB = [
-                (os.path.join(path_to_images_folder, filename), vector)
-                for filename, vector in tmp
-            ]
-            self.LOADED_DB_Name = db_path
-            self.images_folder = path_to_images_folder
+            try:
+                with open(db_path, "rb") as f:
+                    data = orjson.loads(f.read())
 
+                self.MY_LOADED_DB = [
+                    (os.path.join(path_to_images_folder, filename), vector)
+                    for filename, vector in data
+                ]
+                self.LOADED_DB_Name = db_path
+                self.images_folder = path_to_images_folder
+            except Exception as e:
+                print(f"Error loading database file: {e}")
+                return None,
+            
         if img_db is not None:
             self.LOADED_DB = self.MY_LOADED_DB + img_db.LOADED_DB
         else:
@@ -55,6 +73,9 @@ class LoadDB:
         return self, 
 
 class GenerateDB:
+    """
+    Generates a new image database from a folder of images.
+    """    
     def __init__(self):
         pass
     @classmethod
@@ -77,15 +98,27 @@ class GenerateDB:
     RETURN_NAMES = ("ERRORS",)
     OUTPUT_NODE = True
 
-    FUNCTION = "StartGenDB"
+    FUNCTION = "start_gen_db"
     CATEGORY = "ClipVisionTools"
-    def StartGenDB(self, clip_vision, path_to_images_folder, new_db_name):
-        path_to_images = Path(path_to_images_folder)
-        path_to_database = Path(folder_paths.get_folder_paths("EmbDBs")[0] + "\\" + new_db_name)
-        if path_to_database.parent.exists() == False:
-            path_to_database.parent.mkdir()
-        errors = ""
-        if not path_to_database.exists():
-            errors = generate_clip_features_json(clip_vision, path_to_images, path_to_database)
+    def start_gen_db(self, clip_vision, path_to_images_folder, new_db_name):
+        """
+        Generates a new image database.
 
-        return errors,
+        Args:
+            clip_vision: The CLIP vision object.
+            path_to_images_folder: The path to the folder containing the images.
+            new_db_name: The name of the new database file.
+
+        Returns:
+            A string containing any errors that occurred.
+        """        
+        path_to_images = Path(path_to_images_folder)
+        path_to_database = Path(folder_paths.get_folder_paths("EmbDBs")[0]) / new_db_name
+        
+        path_to_database.parent.mkdir(exist_ok=True)
+ 
+        try:
+            errors = generate_clip_features_json(clip_vision, path_to_images, path_to_database)
+            return errors,
+        except Exception as e:
+            return f"Error generating database file: {e}"

image_search.py

@@ -167,10 +167,8 @@ class ResultBrowserAdvanced:
             scores.append(score)
             filenames.append(file_name)
 
-            
         if remove_first_image:
             batch_frame_image = batch_frame_image[1:]
 
-
         return batch_frame_image, filenames, scores