Add files via upload
This commit is contained in:
IMGLAB
@@ -1,60 +1,343 @@
|
||||
# ComfyUI-AutoPickMask (IMGLAB)
|
||||
import { app } from "../../scripts/app.js";
|
||||
|
||||
A ComfyUI custom node that accepts one image and up to 8 mask inputs, compares dimensions, and outputs the first matching mask — with configurable edge-case handling and an optional fallback (else) mask.
|
||||
const MAX_MASKS = 8;
|
||||
const MASK_TYPE = "MASK";
|
||||
const RES_COLOR = "#aaaaaa";
|
||||
const OK_COLOR = "#22cc66";
|
||||
const ERR_COLOR = "#cc3333";
|
||||
const NA_COLOR = "#555555";
|
||||
const ELSE_COLOR = "#cc9900";
|
||||
|
||||
## Installation
|
||||
const SLOT_FONT = "12px Arial";
|
||||
const SLOT_NAME_X = 14;
|
||||
const DOT_RADIUS = 5;
|
||||
const NAME_GAP = 12; // gap: end of slot name → dot centre
|
||||
const DOT_GAP = 10; // gap: dot right edge → resolution text left edge
|
||||
|
||||
Drop the `ComfyUI-AutoPickMask/` folder into `ComfyUI/custom_nodes/` and restart ComfyUI.
|
||||
// LiteGraph draws output labels right-aligned from nodeW-10.
|
||||
// Our widest output name is "mask" or "rgba" (~30px in 12px Arial).
|
||||
// We leave 14px extra buffer, so output area = nodeW - 10 - 30 - 14 = nodeW - 54.
|
||||
// Resolution text right edge must stay at or left of this boundary.
|
||||
const OUTPUT_LABEL_RESERVE = 54;
|
||||
|
||||
```
|
||||
ComfyUI/
|
||||
└── custom_nodes/
|
||||
└── ComfyUI-AutoPickMask/
|
||||
├── __init__.py
|
||||
├── nodes.py
|
||||
├── README.md
|
||||
└── js/
|
||||
└── auto_pick_mask.js
|
||||
```
|
||||
const HELP_SECTIONS = [
|
||||
{
|
||||
title: "Inputs",
|
||||
items: [
|
||||
["image", "Defines the target resolution for matching."],
|
||||
["mask_1..8", "Connect as many masks as needed. A new slot appears as each one is connected."],
|
||||
["else", "Fallback mask. Visible only when no_match is set to Use Else Mask."],
|
||||
]
|
||||
},
|
||||
{
|
||||
title: "Matching",
|
||||
items: [
|
||||
["✓ green", "Mask dimensions match the image exactly."],
|
||||
["✗ red", "Mask dimensions do not match."],
|
||||
["priority", "The lowest-numbered matching slot always wins."],
|
||||
]
|
||||
},
|
||||
{
|
||||
title: "multi_match",
|
||||
items: [
|
||||
["Use First Match", "Silently picks the lowest matching slot."],
|
||||
["Raise Error", "Halts the queue and lists the conflicting slot numbers."],
|
||||
]
|
||||
},
|
||||
{
|
||||
title: "no_match",
|
||||
items: [
|
||||
["Output Empty Mask", "Outputs a zeroed mask (RGBA will be fully transparent)."],
|
||||
["Use Else Mask", "Uses the else input as the fallback output."],
|
||||
["Raise Error", "Halts the queue and lists the sizes that were connected."],
|
||||
]
|
||||
},
|
||||
{
|
||||
title: "Outputs",
|
||||
items: [
|
||||
["image", "Direct passthrough of the input image."],
|
||||
["mask", "The winning mask (or fallback)."],
|
||||
["rgba", "Input image + matched mask as alpha channel [B,H,W,4]."],
|
||||
]
|
||||
},
|
||||
];
|
||||
|
||||
## Inputs
|
||||
const HELP_BTN_W = 16;
|
||||
const HELP_BTN_H = 14;
|
||||
const HELP_BTN_R = 6;
|
||||
|
||||
| Input | Type | Description |
|
||||
|---|---|---|
|
||||
| `image` | IMAGE | Source image — defines the target resolution |
|
||||
| `on_multi_match` | combo | Behaviour when multiple masks match the image size |
|
||||
| `on_no_match` | combo | Behaviour when no mask matches the image size |
|
||||
| `mask_1` … `mask_8` | MASK | Dynamic slots — next slot appears as each one is connected |
|
||||
| `else_mask` | MASK | Fallback mask — only visible when `on_no_match = Use Else Mask` |
|
||||
let _helpPanel = null;
|
||||
|
||||
## Outputs
|
||||
function getHelpPanel() {
|
||||
if (_helpPanel) return _helpPanel;
|
||||
const panel = document.createElement("div");
|
||||
Object.assign(panel.style, {
|
||||
position: "fixed", zIndex: "9999", background: "#1a1a2e",
|
||||
border: "1px solid rgba(255,255,255,0.2)", borderRadius: "8px",
|
||||
padding: "14px 16px", width: "300px", maxHeight: "420px",
|
||||
overflowY: "auto", boxShadow: "0 8px 32px rgba(0,0,0,0.6)",
|
||||
color: "#ddd", fontFamily: "sans-serif", fontSize: "12px",
|
||||
lineHeight: "1.5", display: "none", userSelect: "none",
|
||||
});
|
||||
let html = `<div style="font-size:13px;font-weight:bold;color:#fff;margin-bottom:10px;padding-bottom:6px;border-bottom:1px solid rgba(255,255,255,0.15)">
|
||||
AUTO MASK <span style="color:#a78bfa">(IMGLAB)</span>
|
||||
<span style="font-size:10px;color:#888;font-weight:normal;margin-left:6px">How it works</span>
|
||||
<span id="apm-help-close" style="float:right;cursor:pointer;color:#888;font-size:14px;line-height:1">✕</span>
|
||||
</div>`;
|
||||
for (const section of HELP_SECTIONS) {
|
||||
html += `<div style="margin-top:10px;margin-bottom:4px;font-size:11px;font-weight:bold;color:#8ab4f8;text-transform:uppercase;letter-spacing:0.05em">${section.title}</div>`;
|
||||
html += `<table style="width:100%;border-collapse:collapse">`;
|
||||
for (const [key, val] of section.items)
|
||||
html += `<tr><td style="color:#f0c040;white-space:nowrap;padding:2px 8px 2px 0;vertical-align:top;min-width:110px">${key}</td><td style="color:#ccc;padding:2px 0;vertical-align:top">${val}</td></tr>`;
|
||||
html += `</table>`;
|
||||
}
|
||||
panel.innerHTML = html;
|
||||
document.body.appendChild(panel);
|
||||
panel.querySelector("#apm-help-close").addEventListener("click", () => panel.style.display = "none");
|
||||
document.addEventListener("mousedown", (e) => {
|
||||
if (panel.style.display !== "none" && !panel.contains(e.target)) panel.style.display = "none";
|
||||
}, true);
|
||||
_helpPanel = panel;
|
||||
return panel;
|
||||
}
|
||||
|
||||
| Output | Type | Description |
|
||||
|---|---|---|
|
||||
| `image` | IMAGE | Direct passthrough of the input image |
|
||||
| `matched_mask` | MASK | First mask whose dimensions match the image (or fallback) |
|
||||
| `rgba` | IMAGE | Input image with matched mask as alpha channel `[B,H,W,4]` |
|
||||
function showHelpNear(sx, sy, sw) {
|
||||
const panel = getHelpPanel();
|
||||
panel.style.display = "block";
|
||||
const pw = panel.offsetWidth || 300, ph = panel.offsetHeight || 400;
|
||||
const vw = window.innerWidth, vh = window.innerHeight;
|
||||
let x = sx + sw + 10;
|
||||
if (x + pw > vw - 10) x = sx - pw - 10;
|
||||
let y = sy;
|
||||
if (y + ph > vh - 10) y = vh - ph - 10;
|
||||
if (y < 10) y = 10;
|
||||
panel.style.left = `${x}px`;
|
||||
panel.style.top = `${y}px`;
|
||||
}
|
||||
|
||||
## Edge Case Options
|
||||
app.registerExtension({
|
||||
name: "IMGLAB.AutoMask",
|
||||
|
||||
### on_multi_match — when 2+ masks share the image's exact dimensions:
|
||||
| Option | Behaviour |
|
||||
|---|---|
|
||||
| `Use First (Top Priority)` | Silently picks the lowest-numbered slot |
|
||||
| `Raise Error` | Halts the queue listing the conflicting slot numbers |
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.name !== "AutoMask") return;
|
||||
|
||||
### on_no_match — when no connected mask matches the image size:
|
||||
| Option | Behaviour |
|
||||
|---|---|
|
||||
| `Output Empty Mask` | Outputs a zeroed mask (RGBA will be fully transparent) |
|
||||
| `Use Else Mask` | Reveals an `else_mask` input slot; uses it as the fallback output |
|
||||
| `Raise Error` | Halts the queue listing the sizes that were connected |
|
||||
const origCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
origCreated?.apply(this, arguments);
|
||||
// Default node colour — purple theme
|
||||
if (!this.color) this.color = "#2d1a4a"; // header — dark purple
|
||||
if (!this.bgcolor) this.bgcolor = "#4a2878"; // body — original purple
|
||||
this._am_matchResults = {};
|
||||
this._am_imageSize = null;
|
||||
this._am_maskSizes = {};
|
||||
this._am_elseMaskSize = null;
|
||||
for (let i = this.inputs.length - 1; i >= 0; i--) {
|
||||
const n = this.inputs[i].name;
|
||||
if (/^mask_[2-9]$/.test(n) || /^mask_\d{2,}$/.test(n) || n === "else")
|
||||
this.removeInput(i);
|
||||
}
|
||||
this._am_watchNoMatchWidget();
|
||||
};
|
||||
|
||||
> **Note:** If `Use Else Mask` is selected but `else_mask` is not connected, the node will raise an error rather than silently fail.
|
||||
nodeType.prototype._am_watchNoMatchWidget = function () {
|
||||
const widget = this.widgets?.find(w => w.name === "no_match");
|
||||
if (!widget) return;
|
||||
const self = this;
|
||||
const orig = widget.callback;
|
||||
widget.callback = function (v) { orig?.call(this, v); self._am_syncElseSlot(v); };
|
||||
this._am_syncElseSlot(widget.value);
|
||||
};
|
||||
|
||||
## Node UI
|
||||
nodeType.prototype._am_syncElseSlot = function (value) {
|
||||
const hasElse = this.inputs.some(i => i.name === "else");
|
||||
if (value === "Use Else Mask" && !hasElse) {
|
||||
this.addInput("else", MASK_TYPE);
|
||||
this.setSize(this.computeSize());
|
||||
} else if (value !== "Use Else Mask" && hasElse) {
|
||||
const idx = this.inputs.findIndex(i => i.name === "else");
|
||||
if (idx >= 0) { this.removeInput(idx); this.setSize(this.computeSize()); }
|
||||
}
|
||||
this.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
- **Image slot** — shows resolution (`W×H`) in grey after execution
|
||||
- **Mask slots** — show resolution coloured green (match) or red (mismatch); unconnected slots show a faint grey dot
|
||||
- **Else mask slot** — shows a `fallback` label in amber; only visible when `on_no_match = Use Else Mask`
|
||||
- Mask slots expand automatically as you connect them (up to 8), and collapse when disconnected
|
||||
nodeType.prototype._am_getMaskInputs = function () {
|
||||
return this.inputs
|
||||
.map((inp, idx) => ({ ...inp, _idx: idx }))
|
||||
.filter(inp => /^mask_\d+$/.test(inp.name))
|
||||
.sort((a, b) => parseInt(a.name.split("_")[1]) - parseInt(b.name.split("_")[1]));
|
||||
};
|
||||
|
||||
nodeType.prototype._am_updateDynamicInputs = function () {
|
||||
const masks = this._am_getMaskInputs();
|
||||
let lastConnected = -1;
|
||||
for (let i = 0; i < masks.length; i++) if (masks[i].link != null) lastConnected = i;
|
||||
const targetCount = Math.min(lastConnected + 2, MAX_MASKS);
|
||||
for (let n = masks.length + 1; n <= targetCount; n++) this.addInput(`mask_${n}`, MASK_TYPE);
|
||||
const updated = this._am_getMaskInputs();
|
||||
for (let i = updated.length - 1; i >= targetCount; i--) {
|
||||
const gIdx = this.inputs.findIndex(inp => inp.name === updated[i].name);
|
||||
if (gIdx >= 0 && updated[i].link == null) this.removeInput(gIdx);
|
||||
}
|
||||
// Keep else at the bottom
|
||||
const elseIdx = this.inputs.findIndex(inp => inp.name === "else");
|
||||
if (elseIdx >= 0 && elseIdx !== this.inputs.length - 1) {
|
||||
this.removeInput(elseIdx);
|
||||
this.addInput("else", MASK_TYPE);
|
||||
}
|
||||
this.setSize(this.computeSize());
|
||||
};
|
||||
|
||||
nodeType.prototype.onConnectionsChange = function (type, slotIndex, connected, link_info, input) {
|
||||
if (type !== 1) return;
|
||||
this._am_updateDynamicInputs();
|
||||
this.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
if (message?.match_results?.[0]) {
|
||||
const raw = message.match_results[0];
|
||||
this._am_matchResults = {};
|
||||
for (const k in raw) this._am_matchResults[parseInt(k)] = raw[k];
|
||||
}
|
||||
if (message?.image_size?.[0]) this._am_imageSize = message.image_size[0];
|
||||
if (message?.mask_sizes?.[0]) {
|
||||
const raw = message.mask_sizes[0];
|
||||
this._am_maskSizes = {};
|
||||
for (const k in raw) this._am_maskSizes[parseInt(k)] = raw[k];
|
||||
}
|
||||
this._am_elseMaskSize = message?.else_mask_size?.[0] ?? null;
|
||||
this.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
nodeType.prototype._am_helpBtnRect = function () {
|
||||
const titleH = LiteGraph.NODE_TITLE_HEIGHT;
|
||||
return {
|
||||
x: this.size[0] - HELP_BTN_R - HELP_BTN_W,
|
||||
y: -titleH + (titleH - HELP_BTN_H) / 2,
|
||||
w: HELP_BTN_W, h: HELP_BTN_H,
|
||||
};
|
||||
};
|
||||
|
||||
const origDrawFg = nodeType.prototype.onDrawForeground;
|
||||
nodeType.prototype.onDrawForeground = function (ctx) {
|
||||
origDrawFg?.apply(this, arguments);
|
||||
|
||||
// ? button
|
||||
const btn = this._am_helpBtnRect();
|
||||
ctx.beginPath(); ctx.roundRect(btn.x, btn.y, btn.w, btn.h, 4);
|
||||
ctx.fillStyle = "rgba(255,255,255,0.12)"; ctx.fill();
|
||||
ctx.strokeStyle = "rgba(255,255,255,0.3)"; ctx.lineWidth = 1; ctx.stroke();
|
||||
ctx.font = "bold 11px sans-serif"; ctx.fillStyle = "rgba(255,255,255,0.75)";
|
||||
ctx.textAlign = "center";
|
||||
ctx.fillText("?", btn.x + btn.w / 2, btn.y + btn.h - 3);
|
||||
|
||||
if (this.flags?.collapsed) return;
|
||||
|
||||
const matchResults = this._am_matchResults || {};
|
||||
const maskSizes = this._am_maskSizes || {};
|
||||
const imageSize = this._am_imageSize;
|
||||
const elseMaskSize = this._am_elseMaskSize;
|
||||
const nodeW = this.size[0];
|
||||
|
||||
// Compute shared dot column from widest mask_N name only.
|
||||
// image and else are shorter and would misalign the column.
|
||||
ctx.font = SLOT_FONT;
|
||||
let maxMaskNameW = ctx.measureText("mask_1").width; // minimum
|
||||
for (const inp of this.inputs) {
|
||||
if (/^mask_\d+$/.test(inp.name)) {
|
||||
const w = ctx.measureText(inp.name).width;
|
||||
if (w > maxMaskNameW) maxMaskNameW = w;
|
||||
}
|
||||
}
|
||||
const dotX = SLOT_NAME_X + maxMaskNameW + NAME_GAP + DOT_RADIUS;
|
||||
|
||||
// Resolution text is right-aligned to just before the output label area.
|
||||
// Output labels ("image","mask","rgba") are at the right side of the node.
|
||||
// We reserve OUTPUT_LABEL_RESERVE px from the right for them.
|
||||
const resRight = nodeW - OUTPUT_LABEL_RESERVE;
|
||||
const dimX = dotX+DOT_GAP;
|
||||
|
||||
const drawDot = (cx, cy, color, symbol) => {
|
||||
ctx.beginPath();
|
||||
ctx.arc(cx, cy, DOT_RADIUS, 0, Math.PI * 2);
|
||||
ctx.fillStyle = color; ctx.fill();
|
||||
if (symbol) {
|
||||
ctx.font = "bold 8px sans-serif"; ctx.fillStyle = "#fff";
|
||||
ctx.textAlign = "center";
|
||||
ctx.fillText(symbol, cx, cy + 3);
|
||||
}
|
||||
};
|
||||
|
||||
const drawRes = (text, cy, color) => {
|
||||
ctx.font = "11px monospace"; ctx.fillStyle = color;
|
||||
ctx.textAlign = "left";
|
||||
ctx.fillText(text, dimX, cy + 4);
|
||||
};
|
||||
|
||||
for (let i = 0; i < this.inputs.length; i++) {
|
||||
const inp = this.inputs[i];
|
||||
const canvasPos = this.getConnectionPos(true, i);
|
||||
const cy = canvasPos[1] - this.pos[1];
|
||||
|
||||
// ── image ────────────────────────────────────────────────
|
||||
if (inp.name === "image") {
|
||||
if (imageSize) drawRes(`${imageSize[0]}×${imageSize[1]}`, cy, RES_COLOR);
|
||||
continue;
|
||||
}
|
||||
|
||||
// ── else ─────────────────────────────────────────────────
|
||||
if (inp.name === "else") {
|
||||
ctx.font = SLOT_FONT;
|
||||
const elseNameW = ctx.measureText("else").width;
|
||||
const elseDotX = SLOT_NAME_X + elseNameW + NAME_GAP + DOT_RADIUS;
|
||||
const connected = inp.link != null;
|
||||
drawDot(dotX, cy, connected ? ELSE_COLOR : NA_COLOR, connected && elseMaskSize ? "~" : null);
|
||||
if (connected && elseMaskSize) drawRes(`${elseMaskSize[0]}×${elseMaskSize[1]}`, cy, ELSE_COLOR);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!/^mask_\d+$/.test(inp.name)) continue;
|
||||
|
||||
const num = parseInt(inp.name.split("_")[1]);
|
||||
const connected = inp.link != null;
|
||||
const hasResult = num in matchResults;
|
||||
const hasMaskSz = num in maskSizes;
|
||||
|
||||
// dot at shared column
|
||||
if (connected && hasResult) {
|
||||
const isMatch = matchResults[num];
|
||||
drawDot(dotX, cy, isMatch ? OK_COLOR : ERR_COLOR, isMatch ? "✓" : "✗");
|
||||
} else if (!connected) {
|
||||
ctx.beginPath();
|
||||
ctx.arc(dotX, cy, DOT_RADIUS - 1, 0, Math.PI * 2);
|
||||
ctx.fillStyle = NA_COLOR; ctx.fill();
|
||||
}
|
||||
|
||||
// resolution right-aligned before output labels
|
||||
if (connected && hasMaskSz) {
|
||||
const [mw, mh] = maskSizes[num];
|
||||
const isMatch = hasResult ? matchResults[num] : null;
|
||||
drawRes(`${mw}×${mh}`, cy,
|
||||
isMatch === true ? OK_COLOR : isMatch === false ? ERR_COLOR : RES_COLOR);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const origMouseDown = nodeType.prototype.onMouseDown;
|
||||
nodeType.prototype.onMouseDown = function (e, localPos, graphCanvas) {
|
||||
const btn = this._am_helpBtnRect();
|
||||
const [mx, my] = localPos;
|
||||
if (mx >= btn.x && mx <= btn.x + btn.w && my >= btn.y && my <= btn.y + btn.h) {
|
||||
const canvas = graphCanvas.canvas, rect = canvas.getBoundingClientRect();
|
||||
const scale = graphCanvas.ds?.scale ?? 1, offset = graphCanvas.ds?.offset ?? [0, 0];
|
||||
showHelpNear(
|
||||
rect.left + (this.pos[0] + offset[0]) * scale,
|
||||
rect.top + (this.pos[1] + offset[1]) * scale,
|
||||
this.size[0] * scale
|
||||
);
|
||||
return true;
|
||||
}
|
||||
return origMouseDown?.call(this, e, localPos, graphCanvas);
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user