nodes/advancedtech-sk-ComfyUI-GROUT
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit: ComfyUI-GROUT custom nodes

Browse Source 
Add GROUT segmentation nodes for mosaic grout line detection:
- GROUTSegmentation: Main segmentation node with scale and sliding window modes
- GROUTThreshold: Threshold adjustment without re-running inference
- GROUTHeatmapColorize: Colormap visualization for probability heatmaps

Features:
- Auto-download model from HuggingFace
- Scale mode: resize to 512x512 for inference
- Sliding window mode: process large images with overlapping patches
- Configurable patch size and overlap settings
This commit is contained in:
inventorado7
2026-01-09 10:57:05 +01:00
commit f2fc83e72a
6 changed files with 919 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+109
View File
@@ -0,0 +1,109 @@
# ComfyUI-GROUT
Custom ComfyUI nodes for **GROUT** (Geometric Reasoning Over Unstructured Tessellations) - a deep learning model for detecting and segmenting grout lines in mosaic images.
## Nodes
### GROUT Segmentation
Main segmentation node that processes mosaic images.
**Inputs:**
- `image`: Input mosaic image
- `device`: Select CUDA or CPU for inference
- `threshold`: Initial threshold for binary mask (0.0 - 1.0)
- `model_path` (optional): Custom path to model weights. Leave empty to auto-download from HuggingFace.
**Outputs:**
- `probability_heatmap`: Grayscale probability map (0-1 range)
- `binary_mask`: Thresholded binary mask as image
- `mask`: ComfyUI MASK format for compositing
### GROUT Threshold
Post-processing node to adjust threshold without re-running inference.
**Inputs:**
- `probability_heatmap`: Output from GROUT Segmentation node
- `threshold`: Threshold value (0.0 - 1.0)
**Outputs:**
- `binary_mask`: Thresholded binary mask as image
- `mask`: ComfyUI MASK format
### GROUT Heatmap Colorize
Visualization node to apply colormap to probability heatmap.
**Inputs:**
- `probability_heatmap`: Output from GROUT Segmentation node
- `colormap`: Select from magma, viridis, plasma, inferno, hot, jet
**Outputs:**
- `colored_heatmap`: Colorized visualization
## Installation
### Option 1: ComfyUI Manager
Search for "GROUT" in ComfyUI Manager and install.
### Option 2: Manual Installation
1. Navigate to your ComfyUI custom nodes folder:
```bash
cd ComfyUI/custom_nodes/
```
2. Clone this repository:
```bash
git clone https://github.com/advancedtech-sk/ComfyUI-GROUT.git
```
3. Install dependencies:
```bash
cd ComfyUI-GROUT
pip install -r requirements.txt
```
**For standalone Windows version** (portable/embedded Python):
```bash
cd ComfyUI-GROUT
..\..\..\python_embeded\python.exe -m pip install -r requirements.txt
```
4. Restart ComfyUI
## Model Download
The model will be automatically downloaded from HuggingFace on first use. It will be saved to:
```
ComfyUI/models/grout/grout_b3_zeroshot_v1.pth
```
Alternatively, you can manually download from:
- [HuggingFace Model](https://huggingface.co/advancedtech-sk/GROUT)
## Example Workflow
```
[Load Image] -> [GROUT Segmentation] -> [probability_heatmap] -> [GROUT Threshold] -> [binary_mask]
|
+-> [GROUT Heatmap Colorize] -> [Preview Image]
```
**Workflow Tips:**
1. Use **GROUT Segmentation** to run inference once
2. Connect the `probability_heatmap` output to **GROUT Threshold**
3. Adjust the threshold slider in GROUT Threshold to tune the mask without re-running the model
4. Use **GROUT Heatmap Colorize** for visualization with different colormaps
## Links
- [GROUT Paper (Zenodo)](https://zenodo.org/records/18187265)
- [GROUT Model (HuggingFace)](https://huggingface.co/advancedtech-sk/GROUT)
- [GROUT Demo (HuggingFace Spaces)](https://huggingface.co/spaces/advancedtech-sk/GROUT)
- [GROUT GitHub](https://github.com/advancedtech-sk/GROUT)
## License
MIT License
## Author
Radoslav Lovecky, Institute of Advanced Technologies, Slovakia
Workflows/GROUT_segmentation.json
+371
View File
@@ -0,0 +1,371 @@
{
"id": "394b6835-e282-4241-9fa2-59aabf7df0da",
"revision": 0,
"last_node_id": 8,
"last_link_id": 12,
"nodes": [
{
"id": 7,
"type": "GROUTSegmentation",
"pos": [
768.0085833444832,
784.4327141076791
],
"size": [
270,
218
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 8
}
],
"outputs": [
{
"name": "probability_heatmap",
"type": "IMAGE",
"links": [
9,
11
]
},
{
"name": "binary_mask",
"type": "IMAGE",
"links": [
10
]
},
{
"name": "mask",
"type": "MASK",
"links": null
}
],
"properties": {
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {}
},
"Node name for S&R": "GROUTSegmentation"
},
"widgets_values": [
"cuda",
"sliding_window",
0.5,
512,
64,
""
]
},
{
"id": 3,
"type": "PreviewImage",
"pos": [
1162.2081867641314,
937.443471316358
],
"size": [
140,
246
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 10
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"Node name for S&R": "PreviewImage",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.5.2"
}
},
"widgets_values": []
},
{
"id": 4,
"type": "PreviewImage",
"pos": [
1341.5829607761143,
940.1493486808127
],
"size": [
140,
246
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 12
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"Node name for S&R": "PreviewImage",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.5.2"
}
},
"widgets_values": []
},
{
"id": 8,
"type": "GROUTHeatmapColorize",
"pos": [
1152.6806997969234,
797.2947195049701
],
"size": [
317.595703125,
58
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "probability_heatmap",
"type": "IMAGE",
"link": 11
}
],
"outputs": [
{
"name": "colored_heatmap",
"type": "IMAGE",
"links": [
12
]
}
],
"properties": {
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {}
},
"Node name for S&R": "GROUTHeatmapColorize"
},
"widgets_values": [
"magma"
]
},
{
"id": 5,
"type": "PreviewImage",
"pos": [
1508.9229534605865,
618.2975674422908
],
"size": [
140,
246
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 3
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"Node name for S&R": "PreviewImage",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.5.2"
}
},
"widgets_values": []
},
{
"id": 2,
"type": "GROUTThreshold",
"pos": [
1154.7590602353637,
629.2287088404976
],
"size": [
290.0625,
78
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "probability_heatmap",
"type": "IMAGE",
"link": 9
}
],
"outputs": [
{
"name": "binary_mask",
"type": "IMAGE",
"links": [
3
]
},
{
"name": "mask",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "GROUTThreshold",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.5.2"
}
},
"widgets_values": [
0.7
]
},
{
"id": 6,
"type": "LoadImage",
"pos": [
405.75621728266793,
574.4550289159092
],
"size": [
282.798828125,
314
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
8
]
},
{
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"Node name for S&R": "LoadImage",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.5.2"
}
},
"widgets_values": [
"fish_mosaic.png",
"image"
]
}
],
"links": [
[
3,
2,
0,
5,
0,
"IMAGE"
],
[
8,
6,
0,
7,
0,
"IMAGE"
],
[
9,
7,
0,
2,
0,
"IMAGE"
],
[
10,
7,
1,
3,
0,
"IMAGE"
],
[
11,
7,
0,
8,
0,
"IMAGE"
],
[
12,
8,
0,
4,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ue_links": [],
"links_added_by_ue": [],
"ds": {
"scale": 1.2284597357367253,
"offset": [
-151.1057048390746,
-413.336123673351
]
},
"frontendVersion": "1.34.9",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
Workflows/fish_mosaic.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 9.1 MiB

__init__.py
+17
View File
@@ -0,0 +1,17 @@
"""
ComfyUI-GROUT
Custom nodes for GROUT mosaic grout line segmentation.
GROUT: Geometric Reasoning Over Unstructured Tessellations
Zero-Shot Semantic Segmentation of Mosaic Surfaces via Synthetic Priors
Author: Radoslav Lovecky
License: MIT
"""
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
# Version
__version__ = "1.0.0"
nodes.py
+416
View File
@@ -0,0 +1,416 @@
"""
ComfyUI-GROUT: Custom nodes for GROUT mosaic grout line segmentation
"""
import os
import torch
import numpy as np
from PIL import Image
# Try to import ComfyUI utilities
try:
import folder_paths
except ImportError:
folder_paths = None
class GROUTSegmentation:
"""
GROUT Segmentation Node
Segments grout lines in mosaic images using the GROUT U-Net model.
Outputs both probability heatmap and binary mask.
"""
def __init__(self):
self.model = None
self.current_model_path = None
self.current_device = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"device": (["cuda", "cpu"],),
"mode": (["scale", "sliding_window"],),
"threshold": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"display": "slider"
}),
"patch_size": ("INT", {
"default": 512,
"min": 128,
"max": 1024,
"step": 64,
"tooltip": "Size of patches for sliding window mode"
}),
"overlap": ("INT", {
"default": 64,
"min": 0,
"max": 256,
"step": 16,
"tooltip": "Overlap between patches in sliding window mode"
}),
},
"optional": {
"model_path": ("STRING", {
"default": "",
"multiline": False,
"placeholder": "Leave empty to auto-download from HuggingFace"
}),
}
}
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK")
RETURN_NAMES = ("probability_heatmap", "binary_mask", "mask")
FUNCTION = "segment"
CATEGORY = "GROUT"
def get_default_model_path(self):
"""Get default model path in ComfyUI models directory."""
if folder_paths:
models_dir = folder_paths.models_dir
else:
# Fallback for standalone testing
models_dir = os.path.join(os.path.dirname(__file__), "models")
grout_dir = os.path.join(models_dir, "grout")
os.makedirs(grout_dir, exist_ok=True)
return os.path.join(grout_dir, "grout_b3_zeroshot_v1.pth")
def download_model(self, target_path):
"""Download GROUT model from HuggingFace."""
print("[GROUT] Downloading model from HuggingFace...")
try:
from huggingface_hub import hf_hub_download
hf_hub_download(
repo_id="advancedtech-sk/GROUT",
filename="grout_b3_zeroshot_v1.pth",
local_dir=os.path.dirname(target_path),
local_dir_use_symlinks=False
)
print(f"[GROUT] Model downloaded to: {target_path}")
return True
except ImportError:
print("[GROUT] ERROR: huggingface_hub not installed. Run: pip install huggingface_hub")
return False
except Exception as e:
print(f"[GROUT] ERROR downloading model: {e}")
return False
def load_model(self, model_path, device):
"""Load or reload the GROUT model."""
# Check if we need to reload
if (self.model is not None and
self.current_model_path == model_path and
self.current_device == device):
return self.model
# Determine model path
if not model_path or model_path.strip() == "":
model_path = self.get_default_model_path()
# Download if not exists
if not os.path.exists(model_path):
if not self.download_model(model_path):
raise FileNotFoundError(f"Model not found and download failed: {model_path}")
# Check CUDA availability
if device == "cuda" and not torch.cuda.is_available():
print("[GROUT] CUDA not available, falling back to CPU")
device = "cpu"
print(f"[GROUT] Loading model from: {model_path}")
print(f"[GROUT] Device: {device}")
# Import segmentation_models_pytorch
try:
import segmentation_models_pytorch as smp
except ImportError:
raise ImportError("segmentation_models_pytorch not installed. Run: pip install segmentation-models-pytorch")
# Create model
model = smp.Unet(encoder_name="efficientnet-b3", in_channels=3, classes=1)
# Load weights
checkpoint = torch.load(model_path, map_location=device, weights_only=False)
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
model.load_state_dict(checkpoint['model_state_dict'])
else:
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
self.model = model
self.current_model_path = model_path
self.current_device = device
print("[GROUT] Model loaded successfully")
return model
def segment(self, image, device, mode, threshold, patch_size, overlap, model_path=""):
"""
Segment grout lines in the input image.
Args:
image: ComfyUI IMAGE tensor [B, H, W, C] in range [0, 1]
device: "cuda" or "cpu"
mode: "scale" or "sliding_window"
threshold: Threshold for binary mask (0.0 - 1.0)
patch_size: Size of patches for sliding window mode
overlap: Overlap between patches in sliding window mode
model_path: Optional custom model path
Returns:
probability_heatmap: Grayscale probability map as IMAGE
binary_mask: Thresholded mask as IMAGE
mask: ComfyUI MASK format
"""
# Load model
model = self.load_model(model_path, device)
# Process batch
batch_size = image.shape[0]
prob_maps = []
binary_masks = []
masks = []
# ImageNet normalization
mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3)
std = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3)
for i in range(batch_size):
# Get single image [H, W, C]
img = image[i]
original_h, original_w = img.shape[0], img.shape[1]
if mode == "scale":
prob_map = self._inference_scale(model, img, mean, std, device, original_h, original_w)
else:
prob_map = self._inference_sliding_window(model, img, mean, std, device, original_h, original_w, patch_size, overlap)
# Create probability heatmap as RGB (grayscale) [H, W, 3]
prob_rgb = prob_map.unsqueeze(-1).repeat(1, 1, 3)
prob_maps.append(prob_rgb)
# Create binary mask
binary = (prob_map > threshold).float()
binary_rgb = binary.unsqueeze(-1).repeat(1, 1, 3)
binary_masks.append(binary_rgb)
# MASK format [H, W]
masks.append(binary)
# Stack batches
prob_heatmap = torch.stack(prob_maps, dim=0)
binary_mask_img = torch.stack(binary_masks, dim=0)
mask_out = torch.stack(masks, dim=0)
return (prob_heatmap, binary_mask_img, mask_out)
def _inference_scale(self, model, img, mean, std, device, original_h, original_w):
"""Run inference by scaling image to 512x512."""
# Normalize
img_norm = (img - mean) / std
# Convert to [C, H, W] and add batch dim
img_tensor = img_norm.permute(2, 0, 1).unsqueeze(0).to(device)
# Resize to 512x512 for model
img_resized = torch.nn.functional.interpolate(
img_tensor, size=(512, 512), mode='bilinear', align_corners=False
)
# Inference
with torch.no_grad():
pred = model(img_resized.float())
pred = torch.sigmoid(pred)
# Resize back to original size
pred_resized = torch.nn.functional.interpolate(
pred, size=(original_h, original_w), mode='bilinear', align_corners=False
)
return pred_resized.squeeze().cpu()
def _inference_sliding_window(self, model, img, mean, std, device, original_h, original_w, patch_size, overlap):
"""Run inference using sliding window for large images."""
# If image is smaller than patch size, fall back to scale mode
if original_h <= patch_size and original_w <= patch_size:
return self._inference_scale(model, img, mean, std, device, original_h, original_w)
# Create accumulation tensors
prob_map = torch.zeros((original_h, original_w), dtype=torch.float32)
count_map = torch.zeros((original_h, original_w), dtype=torch.float32)
stride = patch_size - overlap
# Calculate patch positions
y_positions = list(range(0, original_h - patch_size + 1, stride))
if original_h > patch_size and (not y_positions or y_positions[-1] + patch_size < original_h):
y_positions.append(original_h - patch_size)
if not y_positions:
y_positions = [0]
x_positions = list(range(0, original_w - patch_size + 1, stride))
if original_w > patch_size and (not x_positions or x_positions[-1] + patch_size < original_w):
x_positions.append(original_w - patch_size)
if not x_positions:
x_positions = [0]
total_patches = len(y_positions) * len(x_positions)
print(f"[GROUT] Processing {total_patches} patches ({len(y_positions)}x{len(x_positions)})...")
for y in y_positions:
for x in x_positions:
# Extract patch [H, W, C]
patch = img[y:y+patch_size, x:x+patch_size]
# Normalize
patch_norm = (patch - mean) / std
# Convert to [C, H, W] and add batch dim
patch_tensor = patch_norm.permute(2, 0, 1).unsqueeze(0).to(device)
# Inference (patch is already patch_size x patch_size)
with torch.no_grad():
pred = model(patch_tensor.float())
pred = torch.sigmoid(pred)
# Get probability [H, W]
patch_prob = pred.squeeze().cpu()
# Accumulate
prob_map[y:y+patch_size, x:x+patch_size] += patch_prob
count_map[y:y+patch_size, x:x+patch_size] += 1
# Average overlapping regions
prob_map = prob_map / count_map.clamp(min=1)
return prob_map
class GROUTThreshold:
"""
GROUT Threshold Node
Takes a probability heatmap and applies thresholding to create binary mask.
Allows tuning threshold without re-running inference.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"probability_heatmap": ("IMAGE",),
"threshold": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"display": "slider"
}),
},
}
RETURN_TYPES = ("IMAGE", "MASK")
RETURN_NAMES = ("binary_mask", "mask")
FUNCTION = "apply_threshold"
CATEGORY = "GROUT"
def apply_threshold(self, probability_heatmap, threshold):
"""
Apply threshold to probability heatmap.
Args:
probability_heatmap: Grayscale probability map as IMAGE [B, H, W, C]
threshold: Threshold value (0.0 - 1.0)
Returns:
binary_mask: Thresholded mask as IMAGE
mask: ComfyUI MASK format
"""
batch_size = probability_heatmap.shape[0]
binary_masks = []
masks = []
for i in range(batch_size):
# Get probability (use first channel, should be grayscale)
prob = probability_heatmap[i, :, :, 0]
# Apply threshold
binary = (prob > threshold).float()
# Create RGB binary mask
binary_rgb = binary.unsqueeze(-1).repeat(1, 1, 3)
binary_masks.append(binary_rgb)
# MASK format
masks.append(binary)
binary_mask_img = torch.stack(binary_masks, dim=0)
mask_out = torch.stack(masks, dim=0)
return (binary_mask_img, mask_out)
class GROUTHeatmapColorize:
"""
GROUT Heatmap Colorize Node
Converts grayscale probability heatmap to colored visualization.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"probability_heatmap": ("IMAGE",),
"colormap": (["magma", "viridis", "plasma", "inferno", "hot", "jet"],),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("colored_heatmap",)
FUNCTION = "colorize"
CATEGORY = "GROUT"
def colorize(self, probability_heatmap, colormap):
"""
Apply colormap to probability heatmap.
"""
import matplotlib.pyplot as plt
batch_size = probability_heatmap.shape[0]
colored_maps = []
cmap = plt.get_cmap(colormap)
for i in range(batch_size):
# Get probability (grayscale)
prob = probability_heatmap[i, :, :, 0].numpy()
# Apply colormap
colored = cmap(prob)[:, :, :3] # Remove alpha channel
colored_maps.append(torch.from_numpy(colored).float())
colored_output = torch.stack(colored_maps, dim=0)
return (colored_output,)
# Node mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"GROUTSegmentation": GROUTSegmentation,
"GROUTThreshold": GROUTThreshold,
"GROUTHeatmapColorize": GROUTHeatmapColorize,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"GROUTSegmentation": "GROUT Segmentation",
"GROUTThreshold": "GROUT Threshold",
"GROUTHeatmapColorize": "GROUT Heatmap Colorize",
}
requirements.txt
+6
View File
@@ -0,0 +1,6 @@
torch>=2.0.0
segmentation-models-pytorch>=0.3.0
huggingface_hub>=0.20.0
numpy
Pillow
matplotlib