commit f2fc83e72ac571dfdf3e5b16aed536dff98ebaee Author: inventorado7 Date: Fri Jan 9 10:57:05 2026 +0100 Initial commit: ComfyUI-GROUT custom nodes 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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..fc92d36 --- /dev/null +++ b/README.md @@ -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 diff --git a/Workflows/GROUT_segmentation.json b/Workflows/GROUT_segmentation.json new file mode 100644 index 0000000..e37d735 --- /dev/null +++ b/Workflows/GROUT_segmentation.json @@ -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 +} \ No newline at end of file diff --git a/Workflows/fish_mosaic.png b/Workflows/fish_mosaic.png new file mode 100644 index 0000000..39bd21c Binary files /dev/null and b/Workflows/fish_mosaic.png differ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..cbea5c8 --- /dev/null +++ b/__init__.py @@ -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" diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..1b11c77 --- /dev/null +++ b/nodes.py @@ -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", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d8dedee --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +torch>=2.0.0 +segmentation-models-pytorch>=0.3.0 +huggingface_hub>=0.20.0 +numpy +Pillow +matplotlib