nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added image resizing related nodes and a single image prediction node

Browse Source
This commit is contained in:
HowToSD
2025-03-08 03:38:52 -08:00
parent a2fc03f41f
commit 5fee9ffb60
32 changed files with 1007 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/pytorch_wrapper/pt_crop.py
+65
View File
@@ -0,0 +1,65 @@
from typing import Any, Dict
import torch
class PtCrop:
"""
Crops a PyTorch tensor to the specified size. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Crops a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for crop.")
if height > image_height or width > image_width:
raise ValueError("Specified crop size exceeds tensor dimensions.")
# Compute starting position for cropping
top = max(int((height - image_height) / 2), 0)
left = max(int((width - image_width) / 2), 0)
tens2 = tens[..., top:top+height, left:left+width]
return (tens2,)
modules/pytorch_wrapper/pt_from_image.py
+1 -1
View File
@@ -6,7 +6,7 @@ class PtFromImage:
"""
Casts an Image tensor as a PyTorch tensor.
category: PyTorch wrapper - Tensor creation
category: PyTorch wrapper - Image processing
"""
@classmethod
modules/pytorch_wrapper/pt_from_image_transpose.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Any, Dict
import torch
class PtFromImageTranspose:
"""
Casts an image tensor to a PyTorch tensor and transposes it from (H, W, C) to (C, H, W). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"image": ("IMAGE", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, image: torch.Tensor) -> tuple:
"""
Casts an Image tensor as a PyTorch tensor.
Args:
image (torch.Tensor): Image tensor
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
if image.dim() == 3:
tens = image.permute(2, 0, 1)
elif image.dim() == 4:
tens = image.permute(0, 3, 1, 2)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (tens,)
modules/pytorch_wrapper/pt_interpolate_by_scale_factor.py
+56
View File
@@ -0,0 +1,56 @@
from typing import Any, Dict
import torch
class PtInterpolateByScaleFactor:
"""
Resizes a PyTorch tensor using interpolation by scale factor. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"scale_factor": ("FLOAT", {"default":2.0, "min":0.01, "max": 100.0, "step":0.01}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, scale_factor: float, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
scale_factor (float): Scale factor.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, scale_factor=scale_factor, mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_interpolate_to_size.py
+60
View File
@@ -0,0 +1,60 @@
from typing import Any, Dict
import torch
class PtInterpolateToSize:
"""
Resizes a PyTorch tensor using interpolation. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, size=(height, width), mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_load_model.py
+1 -1
View File
@@ -36,6 +36,6 @@ class PtLoadModel:
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
sd = torch.load(model_full_path)
sd = torch.load(model_full_path, weights_only=True)
model.load_state_dict(sd)
return (model,)
modules/pytorch_wrapper/pt_pad.py
+72
View File
@@ -0,0 +1,72 @@
from typing import Any, Dict
import torch
class PtPad:
"""
Pads a PyTorch tensor to the specified size. Padded area will be black. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Pads a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for padding.")
if height < image_height or width < image_width:
raise ValueError("Specified dimensions must be greater than or equal to tensor dimensions.")
# Compute starting position for padding
top = max((height - image_height) // 2, 0)
left = max((width - image_width) // 2, 0)
if dim == 3:
tens2 = torch.zeros(tens.size(0), height, width,
dtype=tens.dtype, device=tens.device)
else:
tens2 = torch.zeros(tens.size(0), tens.size(1), height, width,
dtype=tens.dtype, device=tens.device)
tens2[..., top:top+image_height, left:left+image_width] = tens
return (tens2,)
modules/pytorch_wrapper/pt_predict_classification_model.py
+66
View File
@@ -0,0 +1,66 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtPredictClassificationModel:
"""
Performs inference on input data.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"inputs": ("TENSOR", {}),
"class_id_to_name_map": ("PYDICT", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("class_name", "class_id", "probability")
RETURN_TYPES: Tuple[str, ...] = ("STRING", "INT", "FLOAT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module,
inputs: torch.Tensor,
class_id_to_name_map: Dict,
use_gpu: bool) -> Tuple[str, int, float]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
inputs (torch.Tensor): Input image to predict the class for.
class_id_to_name_map (Dict): A Python dictionary contains the mapping of
class ID to class name.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[str, int, float]: A tuple containing "class_name", "class_id", "probability".
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
x = torch.unsqueeze(inputs, 0)
if use_gpu:
x = x.to("cuda")
outputs = model(x)
class_id = outputs.argmax(dim=1)[0].detach().cpu().item()
prob = F.softmax(outputs, dim=1)[0][class_id].detach().cpu().item()
class_name = class_id_to_name_map[class_id]
return (class_name, class_id, prob)
modules/pytorch_wrapper/pt_to_image.py
+2 -1
View File
@@ -6,7 +6,7 @@ class PtToImage:
"""
Casts a PyTorch tensor as an Image tensor.
category: PyTorch wrapper - Tensor data conversion
category: PyTorch wrapper - Image processing
"""
@classmethod
@@ -23,6 +23,7 @@ class PtToImage:
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
modules/pytorch_wrapper/pt_to_image_transpose.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Any, Dict
import torch
class PtToImageTranspose:
"""
Casts a PyTorch tensor as an Image tensor and transposes it from (C, H, W) to (H, W, C). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Casts a PyTorch tensor as an Image tensor.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the image tensor.
"""
if tens.dim() == 3:
image = tens.permute(1, 2, 0)
elif tens.dim() == 4:
image = tens.permute(0, 2, 3, 1)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (image,)
modules/pytorch_wrapper/pto_adam.py
+2 -2
View File
@@ -5,9 +5,9 @@ import torch.nn.functional as F
class PtoAdam:
"""
Converts PyTorch tensor to NumPy ndarray.
Instantiates the Adam optimizer.
category: PyTorch wrapper - Tensor data conversion
category: PyTorch wrapper - Training
"""
@classmethod