116 lines
3.7 KiB
Python
116 lines
3.7 KiB
Python
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
import folder_paths
|
|
|
|
|
|
def _tensor_image_to_pil_image(tensor):
|
|
image_tensor = tensor[0]
|
|
if hasattr(image_tensor, "mul") and hasattr(image_tensor, "clamp"):
|
|
image_array = image_tensor.mul(255).clamp(0, 255)
|
|
if hasattr(image_array, "byte"):
|
|
image_array = image_array.byte()
|
|
image_array = image_array.cpu().numpy()
|
|
return Image.fromarray(image_array)
|
|
|
|
image_array = 255.0 * image_tensor.cpu().numpy()
|
|
return Image.fromarray(np.clip(image_array, 0, 255).astype(np.uint8))
|
|
|
|
|
|
class DumasImageCompareNode:
|
|
DESCRIPTION = (
|
|
"Dumas Image Compare shows the difference between two images directly on "
|
|
"the node. Connect one or two IMAGE inputs to compare before/after "
|
|
"results, model variants, or processing stages without breaking a "
|
|
"workflow when one branch is bypassed."
|
|
)
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("new image",)
|
|
FUNCTION = "compare_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Dumas/Image"
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
self.prefix_append = "_dumascmp_" + "".join(
|
|
random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)
|
|
)
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"optional": {
|
|
"image1": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"First image to compare. Optional so muted or bypassed "
|
|
"branches do not trigger a missing-input error."
|
|
)
|
|
},
|
|
),
|
|
"image2": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"Second image to compare. Optional so the node can still "
|
|
"display a single available image."
|
|
)
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
def compare_images(self, image1=None, image2=None):
|
|
present = []
|
|
if image1 is not None:
|
|
present.append((1, image1))
|
|
if image2 is not None:
|
|
present.append((2, image2))
|
|
results = []
|
|
|
|
if present:
|
|
first_tensor = present[0][1]
|
|
prefix = "dumas_compare" + self.prefix_append
|
|
first_image = first_tensor[0]
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
prefix,
|
|
self.output_dir,
|
|
first_image.shape[1],
|
|
first_image.shape[0],
|
|
)
|
|
join_path = os.path.join
|
|
for slot, tensor in present:
|
|
image = _tensor_image_to_pil_image(tensor)
|
|
file_name = f"{filename}_{counter:05}_.png"
|
|
image.save(
|
|
join_path(full_output_folder, file_name),
|
|
compress_level=self.compress_level,
|
|
)
|
|
results.append(
|
|
{
|
|
"filename": file_name,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
"slot": slot,
|
|
}
|
|
)
|
|
counter += 1
|
|
|
|
new_image = image2 if image2 is not None else image1
|
|
return {"ui": {"images": results}, "result": (new_image,)}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DumasImageCompare": DumasImageCompareNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DumasImageCompare": "Dumas Image Compare",
|
|
}
|