diff --git a/README.md b/README.md index 103b661..ed7db3a 100644 --- a/README.md +++ b/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) diff --git a/__init__.py b/__init__.py index 6acd042..95ca02f 100644 --- a/__init__.py +++ b/__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" } diff --git a/clipvision_db.py b/clipvision_db.py index 099b2d7..a96b732 100644 --- a/clipvision_db.py +++ b/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}" diff --git a/image_search.py b/image_search.py index a1abe8b..8b4577d 100644 --- a/image_search.py +++ b/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