Mask Refine (K-Means)

The source file is blunt about why this node exists: "Helps grow under-segmented masks from SAM3 or other segmentation models." If you've used SAM-class segmenters, you know the failure - the mask stops short, leaving the subject's edges clipped, and a blind dilation just swallows background along with the subject. Mask Refine (K-Means) grows a mask into the image instead: it looks at the colors around the mask edge and includes only pixels that are similar to the subject you already selected.

That's a genuinely different operation from plain dilation, and it's the one people are actually asking for when they hit the under-segmentation wall. The object is a person against a wall - the mask clips the hair or the jacket. Dilation grabs wall. This node grabs more person. It's also the rare node where the author bothered to write real tooltips on every input, which tells you it's meant to be tuned, not just dropped in.

How it works

The algorithm is textbook and effective. It takes the image plus your rough mask, dilates the mask outward by up to grow_pixels to define a search region, then runs k-means clustering (n_clusters groups) over the pixel colors in that region. It finds the cluster closest to the average color of your original mask - that's "the subject's color" - and keeps only the boundary pixels that belong to that cluster and are within similarity_threshold of the object color. The result is then passed through morphological open/close and a distance-transform smoothing pass, which is how you get a clean edge instead of a confetti outline. Batches are handled per-image.

Inputs and outputs

The three knobs, all with the author's own descriptions:

• n_clusters - 2 to 10, default 3. More clusters capture more color variation but run slower. For a subject with a few distinct tones, 3 is a fine start
• grow_pixels - 0 to 200, default 50. How far the mask is allowed to spread
• similarity_threshold - 0 to 1, default 0.5. Higher = stricter, only very-similar colors get included. Lower = greedier growth

Inputs are image and mask; the output is refined_mask, a MASK.

Installing it

Standard XJNodes install. ComfyUI Manager → search "ComfyUI-XJNodes" → install → restart. Or:

cd ComfyUI/custom_nodes
git clone https://github.com/alexjx/ComfyUI-XJNodes
# restart ComfyUI

This is the one XJNodes node with a real extra dependency. It imports sklearn.cluster.KMeans, and scikit-learn is not guaranteed in a ComfyUI install. If you hit ModuleNotFoundError: No module named 'sklearn', run pip install scikit-learn and restart. Category: XJNodes/segs.

Tuning advice, earned from the failure modes

If the refined mask is too small, the subject's colors span more clusters than you guessed - raise n_clusters or loosen similarity_threshold. If it's bleeding into the background, the background shares too much color with the subject - raise similarity_threshold and lower grow_pixels so the search region doesn't reach the trouble. If it barely changes at all, grow_pixels is the limit you're hitting, not the threshold. It's a power tool with three controls that do exactly what their tooltips say; spend one run with a Preview wired to the output and you'll have it dialed in. For a lighter-weight alternative that uses edge detection instead of clustering, the pack's Mask Refine (Morph) is the other half of this pair.