inventorado7
f2fc83e72a
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
417 lines
14 KiB
Python
417 lines
14 KiB
Python
"""
|
|
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",
|
|
}
|