mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-21 20:52:11 -03:00
100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
|
|
class ImageDetails:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_input": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT", "BOOL", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("WIDTH", "HEIGHT", "HAS_TRANSPARENCY", "ORIENTATION", "TYPE", "ALL")
|
|
FUNCTION = "show_image_details"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Bjornulf"
|
|
|
|
def show_image_details(self, image_input):
|
|
if isinstance(image_input, torch.Tensor):
|
|
is_tensor = True
|
|
input_type = "tensor"
|
|
# Ensure the tensor is on CPU and convert to numpy
|
|
image_input = image_input.cpu().numpy()
|
|
elif isinstance(image_input, (bytes, bytearray)):
|
|
is_tensor = False
|
|
input_type = "bytes"
|
|
image_input = [image_input] # Wrap single bytes object in a list
|
|
else:
|
|
is_tensor = False
|
|
input_type = "bytes"
|
|
|
|
all_widths, all_heights, all_transparencies, all_details, all_orientations = [], [], [], [], []
|
|
|
|
if is_tensor:
|
|
# Handle tensor images
|
|
if len(image_input.shape) == 5: # (batch, 1, channels, height, width)
|
|
image_input = np.squeeze(image_input, axis=1)
|
|
|
|
batch_size = image_input.shape[0]
|
|
for i in range(batch_size):
|
|
image = image_input[i]
|
|
|
|
# Ensure the image is in HxWxC format
|
|
if image.shape[0] == 3 or image.shape[0] == 4: # If it's in CxHxW format
|
|
image = np.transpose(image, (1, 2, 0)) # Change to HxWxC
|
|
|
|
# Normalize to 0-255 range if necessary
|
|
if image.max() <= 1:
|
|
image = (image * 255).astype('uint8')
|
|
else:
|
|
image = image.astype('uint8')
|
|
|
|
pil_image = Image.fromarray(image)
|
|
self.process_image(pil_image, input_type, all_widths, all_heights, all_transparencies, all_details, all_orientations)
|
|
else:
|
|
# Handle bytes-like objects
|
|
batch_size = len(image_input)
|
|
for i in range(batch_size):
|
|
pil_image = Image.open(io.BytesIO(image_input[i]))
|
|
self.process_image(pil_image, input_type, all_widths, all_heights, all_transparencies, all_details, all_orientations)
|
|
|
|
# Combine all details into a single string
|
|
combined_details = "\n".join(all_details)
|
|
|
|
# Return the details of the first image, plus the combined details string
|
|
return (all_widths[0], all_heights[0], all_transparencies[0], all_orientations[0],
|
|
input_type, combined_details)
|
|
|
|
def process_image(self, pil_image, input_type, all_widths, all_heights, all_transparencies, all_details, all_orientations):
|
|
# Get image details
|
|
width, height = pil_image.size
|
|
has_transparency = pil_image.mode in ('RGBA', 'LA') or \
|
|
(pil_image.mode == 'P' and 'transparency' in pil_image.info)
|
|
|
|
# Determine orientation
|
|
if width > height:
|
|
orientation = "landscape"
|
|
elif height > width:
|
|
orientation = "portrait"
|
|
else:
|
|
orientation = "square"
|
|
|
|
# Prepare the ALL string
|
|
details = f"\nType: {input_type}"
|
|
details += f"\nWidth: {width}"
|
|
details += f"\nHeight: {height}"
|
|
details += f"\nNumber of Pixels: {width * height}"
|
|
details += f"\nLoaded with transparency: {has_transparency}"
|
|
details += f"\nImage Mode: {pil_image.mode}"
|
|
details += f"\nOrientation: {orientation}\n"
|
|
|
|
all_widths.append(width)
|
|
all_heights.append(height)
|
|
all_transparencies.append(has_transparency)
|
|
all_details.append(details)
|
|
all_orientations.append(orientation)
|