536 lines
19 KiB
JavaScript
536 lines
19 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
|
|
function chainCallback(object, property, callback) {
|
|
const original = object[property];
|
|
if (original) {
|
|
object[property] = function (...args) {
|
|
const result = original.apply(this, args);
|
|
callback.apply(this, args);
|
|
return result;
|
|
};
|
|
} else {
|
|
object[property] = callback;
|
|
}
|
|
}
|
|
|
|
function findWidget(node, name) {
|
|
return node.widgets?.find((widget) => widget.name === name);
|
|
}
|
|
|
|
function clamp(value, min, max) {
|
|
return Math.max(min, Math.min(max, value));
|
|
}
|
|
|
|
function parseRegions(value) {
|
|
try {
|
|
const parsed = JSON.parse(value || "[]");
|
|
if (Array.isArray(parsed)) return parsed;
|
|
if (parsed && Array.isArray(parsed.regions)) return parsed.regions;
|
|
if (parsed && typeof parsed === "object") return [parsed];
|
|
} catch (error) {
|
|
console.warn("UR Region Editor: failed to parse regions_json", error);
|
|
}
|
|
return [];
|
|
}
|
|
|
|
function normalizeRect(region) {
|
|
const x = Number(region.x ?? region.x1 ?? 0);
|
|
const y = Number(region.y ?? region.y1 ?? 0);
|
|
if (region.w != null || region.h != null) {
|
|
return {
|
|
x,
|
|
y,
|
|
w: Number(region.w ?? 0),
|
|
h: Number(region.h ?? 0),
|
|
};
|
|
}
|
|
const x2 = Number(region.x2 ?? region.right ?? x);
|
|
const y2 = Number(region.y2 ?? region.bottom ?? y);
|
|
return {
|
|
x: Math.min(x, x2),
|
|
y: Math.min(y, y2),
|
|
w: Math.abs(x2 - x),
|
|
h: Math.abs(y2 - y),
|
|
};
|
|
}
|
|
|
|
function createButton(label, title) {
|
|
const button = document.createElement("button");
|
|
button.textContent = label;
|
|
button.title = title || label;
|
|
button.style.cssText = "height:24px;padding:0 8px;border:1px solid #555;background:#2b2b2b;color:#ddd;border-radius:4px;cursor:pointer;font-size:12px;";
|
|
button.addEventListener("mouseenter", () => {
|
|
button.style.background = "#3a3a3a";
|
|
});
|
|
button.addEventListener("mouseleave", () => {
|
|
button.style.background = button.dataset.active === "true" ? "#4b5f7a" : "#2b2b2b";
|
|
});
|
|
return button;
|
|
}
|
|
|
|
function createNumberInput(value, min, title) {
|
|
const input = document.createElement("input");
|
|
input.type = "number";
|
|
input.value = String(value);
|
|
input.min = String(min ?? 0);
|
|
input.title = title || "";
|
|
input.style.cssText = "height:22px;width:76px;background:#111;color:#ddd;border:1px solid #555;border-radius:4px;padding:0 4px;font-size:12px;";
|
|
return input;
|
|
}
|
|
|
|
class RegionEditorWidget {
|
|
constructor(node) {
|
|
this.node = node;
|
|
this.regionsWidget = findWidget(node, "regions_json");
|
|
this.widthWidget = findWidget(node, "canvas_width");
|
|
this.heightWidget = findWidget(node, "canvas_height");
|
|
this.defaultFrameWidget = findWidget(node, "default_frame_index");
|
|
|
|
this.regions = parseRegions(this.regionsWidget?.value);
|
|
this.mode = "rect";
|
|
this.activeFrame = Number(this.defaultFrameWidget?.value ?? 0);
|
|
this.selectedIndex = -1;
|
|
this.dragStart = null;
|
|
this.currentDraft = null;
|
|
this.polygonDraft = [];
|
|
this.background = null;
|
|
this.video = null;
|
|
this.videoFps = 24;
|
|
|
|
this.root = document.createElement("div");
|
|
this.root.style.cssText = "display:flex;flex-direction:column;gap:6px;width:100%;height:100%;box-sizing:border-box;background:#121212;border:1px solid #303030;border-radius:6px;padding:6px;overflow:hidden;";
|
|
|
|
this.toolbar = document.createElement("div");
|
|
this.toolbar.style.cssText = "display:flex;align-items:center;gap:6px;flex-wrap:wrap;";
|
|
|
|
this.rectBtn = createButton("Rect", "Draw rectangle by dragging");
|
|
this.circleBtn = createButton("Circle", "Draw circle by dragging");
|
|
this.polyBtn = createButton("Poly", "Click points, then Finish Poly");
|
|
this.finishPolyBtn = createButton("Finish Poly", "Finish current polygon");
|
|
this.undoPointBtn = createButton("Undo Point", "Remove last polygon point");
|
|
this.deleteBtn = createButton("Delete", "Delete selected region");
|
|
this.clearBtn = createButton("Clear", "Clear all regions");
|
|
this.loadBgBtn = createButton("Load BG", "Load local image/video as editor background only");
|
|
|
|
this.frameInput = createNumberInput(this.activeFrame, 0, "Current keyframe index");
|
|
this.frameLabel = document.createElement("span");
|
|
this.frameLabel.textContent = "Frame";
|
|
this.frameLabel.style.cssText = "color:#aaa;font-size:12px;";
|
|
|
|
this.fileInput = document.createElement("input");
|
|
this.fileInput.type = "file";
|
|
this.fileInput.accept = "image/*,video/*";
|
|
this.fileInput.style.display = "none";
|
|
|
|
this.videoControls = document.createElement("div");
|
|
this.videoControls.style.cssText = "display:none;align-items:center;gap:6px;width:100%;";
|
|
this.playBtn = createButton("Play", "Play/pause background video");
|
|
this.videoSlider = document.createElement("input");
|
|
this.videoSlider.type = "range";
|
|
this.videoSlider.min = "0";
|
|
this.videoSlider.max = "0";
|
|
this.videoSlider.step = "0.001";
|
|
this.videoSlider.value = "0";
|
|
this.videoSlider.style.cssText = "flex:1;";
|
|
this.fpsInput = createNumberInput(this.videoFps, 1, "FPS used to convert video time to frame index");
|
|
this.useVideoFrameBtn = createButton("Use Time", "Set keyframe from current video time");
|
|
|
|
this.canvas = document.createElement("canvas");
|
|
this.canvas.width = Number(this.widthWidget?.value ?? 768);
|
|
this.canvas.height = Number(this.heightWidget?.value ?? 512);
|
|
this.canvas.style.cssText = "width:100%;height:auto;max-height:520px;object-fit:contain;background:#181818;border:1px solid #333;border-radius:4px;cursor:crosshair;";
|
|
this.ctx = this.canvas.getContext("2d");
|
|
|
|
this.status = document.createElement("div");
|
|
this.status.style.cssText = "color:#aaa;font-size:12px;white-space:pre-wrap;line-height:1.3;";
|
|
|
|
this.toolbar.append(
|
|
this.rectBtn,
|
|
this.circleBtn,
|
|
this.polyBtn,
|
|
this.finishPolyBtn,
|
|
this.undoPointBtn,
|
|
this.deleteBtn,
|
|
this.clearBtn,
|
|
this.loadBgBtn,
|
|
this.frameLabel,
|
|
this.frameInput,
|
|
this.fileInput,
|
|
);
|
|
this.videoControls.append(this.playBtn, this.videoSlider, document.createTextNode("FPS"), this.fpsInput, this.useVideoFrameBtn);
|
|
this.root.append(this.toolbar, this.videoControls, this.canvas, this.status);
|
|
|
|
this.bindEvents();
|
|
this.setMode("rect");
|
|
this.syncWidgets(false);
|
|
this.render();
|
|
}
|
|
|
|
bindEvents() {
|
|
this.rectBtn.onclick = () => this.setMode("rect");
|
|
this.circleBtn.onclick = () => this.setMode("circle");
|
|
this.polyBtn.onclick = () => this.setMode("polygon");
|
|
this.finishPolyBtn.onclick = () => this.finishPolygon();
|
|
this.undoPointBtn.onclick = () => {
|
|
this.polygonDraft.pop();
|
|
this.render();
|
|
};
|
|
this.deleteBtn.onclick = () => this.deleteSelected();
|
|
this.clearBtn.onclick = () => {
|
|
this.regions = [];
|
|
this.selectedIndex = -1;
|
|
this.polygonDraft = [];
|
|
this.syncWidgets(true);
|
|
this.render();
|
|
};
|
|
this.loadBgBtn.onclick = () => this.fileInput.click();
|
|
|
|
this.frameInput.addEventListener("change", () => {
|
|
this.activeFrame = Math.max(0, Number(this.frameInput.value || 0));
|
|
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
|
this.render();
|
|
});
|
|
this.fpsInput.addEventListener("change", () => {
|
|
this.videoFps = Math.max(1, Number(this.fpsInput.value || 24));
|
|
});
|
|
|
|
this.fileInput.addEventListener("change", () => this.loadBackgroundFile());
|
|
this.playBtn.onclick = () => this.toggleVideo();
|
|
this.videoSlider.addEventListener("input", () => {
|
|
if (!this.video) return;
|
|
this.video.currentTime = Number(this.videoSlider.value || 0);
|
|
this.drawVideoFrame();
|
|
});
|
|
this.useVideoFrameBtn.onclick = () => {
|
|
if (!this.video) return;
|
|
this.activeFrame = Math.max(0, Math.round(this.video.currentTime * this.videoFps));
|
|
this.frameInput.value = String(this.activeFrame);
|
|
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
|
this.render();
|
|
};
|
|
|
|
this.canvas.addEventListener("mousedown", (event) => this.onMouseDown(event));
|
|
this.canvas.addEventListener("mousemove", (event) => this.onMouseMove(event));
|
|
this.canvas.addEventListener("mouseup", (event) => this.onMouseUp(event));
|
|
this.canvas.addEventListener("dblclick", () => this.finishPolygon());
|
|
this.canvas.addEventListener("contextmenu", (event) => {
|
|
event.preventDefault();
|
|
const point = this.toCanvasPoint(event);
|
|
const index = this.hitTest(point.x, point.y);
|
|
if (index >= 0) {
|
|
this.selectedIndex = index;
|
|
this.deleteSelected();
|
|
}
|
|
});
|
|
document.addEventListener("keydown", (event) => {
|
|
if (!this.root.isConnected) return;
|
|
if (event.key === "Delete" || event.key === "Backspace") this.deleteSelected();
|
|
if (event.key === "Enter" && this.mode === "polygon") this.finishPolygon();
|
|
if (event.key === "Escape") {
|
|
this.currentDraft = null;
|
|
this.polygonDraft = [];
|
|
this.render();
|
|
}
|
|
});
|
|
|
|
if (this.regionsWidget) {
|
|
const previous = this.regionsWidget.callback;
|
|
this.regionsWidget.callback = (...args) => {
|
|
if (previous) previous.apply(this.regionsWidget, args);
|
|
this.regions = parseRegions(this.regionsWidget.value);
|
|
this.render();
|
|
};
|
|
}
|
|
}
|
|
|
|
setMode(mode) {
|
|
this.mode = mode;
|
|
for (const [button, name] of [[this.rectBtn, "rect"], [this.circleBtn, "circle"], [this.polyBtn, "polygon"]]) {
|
|
button.dataset.active = name === mode ? "true" : "false";
|
|
button.style.background = name === mode ? "#4b5f7a" : "#2b2b2b";
|
|
}
|
|
this.currentDraft = null;
|
|
this.render();
|
|
}
|
|
|
|
syncWidgets(markDirty) {
|
|
const text = JSON.stringify(this.regions);
|
|
if (this.regionsWidget) this.regionsWidget.value = text;
|
|
if (this.widthWidget) this.widthWidget.value = this.canvas.width;
|
|
if (this.heightWidget) this.heightWidget.value = this.canvas.height;
|
|
if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame;
|
|
if (markDirty) app.graph?.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
loadBackgroundFile() {
|
|
const file = this.fileInput.files?.[0];
|
|
if (!file) return;
|
|
const url = URL.createObjectURL(file);
|
|
this.background = null;
|
|
this.video = null;
|
|
this.videoControls.style.display = "none";
|
|
if (file.type.startsWith("video/")) {
|
|
const video = document.createElement("video");
|
|
video.src = url;
|
|
video.muted = true;
|
|
video.playsInline = true;
|
|
video.addEventListener("loadedmetadata", () => {
|
|
this.video = video;
|
|
this.canvas.width = video.videoWidth || this.canvas.width;
|
|
this.canvas.height = video.videoHeight || this.canvas.height;
|
|
this.videoSlider.max = String(video.duration || 0);
|
|
this.videoControls.style.display = "flex";
|
|
this.drawVideoFrame();
|
|
this.syncWidgets(true);
|
|
});
|
|
video.addEventListener("timeupdate", () => {
|
|
this.videoSlider.value = String(video.currentTime);
|
|
this.drawVideoFrame();
|
|
});
|
|
return;
|
|
}
|
|
|
|
const image = new Image();
|
|
image.onload = () => {
|
|
this.background = image;
|
|
this.canvas.width = image.width;
|
|
this.canvas.height = image.height;
|
|
this.syncWidgets(true);
|
|
this.render();
|
|
};
|
|
image.src = url;
|
|
}
|
|
|
|
toggleVideo() {
|
|
if (!this.video) return;
|
|
if (this.video.paused) {
|
|
this.video.play();
|
|
this.playBtn.textContent = "Pause";
|
|
} else {
|
|
this.video.pause();
|
|
this.playBtn.textContent = "Play";
|
|
}
|
|
}
|
|
|
|
drawVideoFrame() {
|
|
this.render();
|
|
}
|
|
|
|
toCanvasPoint(event) {
|
|
const rect = this.canvas.getBoundingClientRect();
|
|
return {
|
|
x: clamp((event.clientX - rect.left) * (this.canvas.width / rect.width), 0, this.canvas.width),
|
|
y: clamp((event.clientY - rect.top) * (this.canvas.height / rect.height), 0, this.canvas.height),
|
|
};
|
|
}
|
|
|
|
onMouseDown(event) {
|
|
event.preventDefault();
|
|
const point = this.toCanvasPoint(event);
|
|
const hit = this.hitTest(point.x, point.y);
|
|
if (hit >= 0 && event.shiftKey) {
|
|
this.selectedIndex = hit;
|
|
this.render();
|
|
return;
|
|
}
|
|
if (this.mode === "polygon") {
|
|
this.polygonDraft.push(point);
|
|
this.render();
|
|
return;
|
|
}
|
|
this.dragStart = point;
|
|
this.currentDraft = { type: this.mode, frame_index: this.activeFrame, start: point, end: point };
|
|
}
|
|
|
|
onMouseMove(event) {
|
|
if (!this.currentDraft) return;
|
|
this.currentDraft.end = this.toCanvasPoint(event);
|
|
this.render();
|
|
}
|
|
|
|
onMouseUp(event) {
|
|
if (!this.currentDraft) return;
|
|
this.currentDraft.end = this.toCanvasPoint(event);
|
|
const region = this.draftToRegion(this.currentDraft);
|
|
this.currentDraft = null;
|
|
if (region) {
|
|
this.regions.push(region);
|
|
this.selectedIndex = this.regions.length - 1;
|
|
this.syncWidgets(true);
|
|
}
|
|
this.render();
|
|
}
|
|
|
|
draftToRegion(draft) {
|
|
const x1 = draft.start.x;
|
|
const y1 = draft.start.y;
|
|
const x2 = draft.end.x;
|
|
const y2 = draft.end.y;
|
|
if (draft.type === "rect") {
|
|
const x = Math.min(x1, x2);
|
|
const y = Math.min(y1, y2);
|
|
const w = Math.abs(x2 - x1);
|
|
const h = Math.abs(y2 - y1);
|
|
if (w < 3 || h < 3) return null;
|
|
return { type: "rect", frame_index: this.activeFrame, x: Math.round(x), y: Math.round(y), w: Math.round(w), h: Math.round(h) };
|
|
}
|
|
if (draft.type === "circle") {
|
|
const r = Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2);
|
|
if (r < 3) return null;
|
|
return { type: "circle", frame_index: this.activeFrame, cx: Math.round(x1), cy: Math.round(y1), r: Math.round(r) };
|
|
}
|
|
return null;
|
|
}
|
|
|
|
finishPolygon() {
|
|
if (this.polygonDraft.length < 3) return;
|
|
this.regions.push({
|
|
type: "polygon",
|
|
frame_index: this.activeFrame,
|
|
points: this.polygonDraft.map((point) => [Math.round(point.x), Math.round(point.y)]),
|
|
});
|
|
this.selectedIndex = this.regions.length - 1;
|
|
this.polygonDraft = [];
|
|
this.syncWidgets(true);
|
|
this.render();
|
|
}
|
|
|
|
deleteSelected() {
|
|
if (this.selectedIndex < 0 || this.selectedIndex >= this.regions.length) return;
|
|
this.regions.splice(this.selectedIndex, 1);
|
|
this.selectedIndex = -1;
|
|
this.syncWidgets(true);
|
|
this.render();
|
|
}
|
|
|
|
hitTest(x, y) {
|
|
for (let i = this.regions.length - 1; i >= 0; i--) {
|
|
const region = this.regions[i];
|
|
if (Number(region.frame_index ?? 0) !== this.activeFrame) continue;
|
|
const type = String(region.type || "rect").toLowerCase();
|
|
if (type === "rect") {
|
|
const rect = normalizeRect(region);
|
|
if (x >= rect.x && x <= rect.x + rect.w && y >= rect.y && y <= rect.y + rect.h) return i;
|
|
} else if (type === "circle") {
|
|
const dx = x - Number(region.cx ?? region.x ?? 0);
|
|
const dy = y - Number(region.cy ?? region.y ?? 0);
|
|
if (dx * dx + dy * dy <= Number(region.r ?? region.radius ?? 0) ** 2) return i;
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
render() {
|
|
const ctx = this.ctx;
|
|
ctx.clearRect(0, 0, this.canvas.width, this.canvas.height);
|
|
ctx.fillStyle = "#181818";
|
|
ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
|
|
if (this.video) ctx.drawImage(this.video, 0, 0, this.canvas.width, this.canvas.height);
|
|
else if (this.background) ctx.drawImage(this.background, 0, 0, this.canvas.width, this.canvas.height);
|
|
else this.drawGrid();
|
|
|
|
this.regions.forEach((region, index) => this.drawRegion(region, index === this.selectedIndex));
|
|
if (this.currentDraft) this.drawDraft(this.currentDraft);
|
|
if (this.polygonDraft.length) this.drawPolygonDraft();
|
|
this.status.textContent = `Mode: ${this.mode} | Frame: ${this.activeFrame} | Regions: ${this.regions.length}\nLoad BG is editor-only; regions_json is portable and saved in the workflow.`;
|
|
}
|
|
|
|
drawGrid() {
|
|
const ctx = this.ctx;
|
|
ctx.strokeStyle = "#2a2a2a";
|
|
ctx.lineWidth = 1;
|
|
for (let x = 0; x <= this.canvas.width; x += 64) {
|
|
ctx.beginPath();
|
|
ctx.moveTo(x, 0);
|
|
ctx.lineTo(x, this.canvas.height);
|
|
ctx.stroke();
|
|
}
|
|
for (let y = 0; y <= this.canvas.height; y += 64) {
|
|
ctx.beginPath();
|
|
ctx.moveTo(0, y);
|
|
ctx.lineTo(this.canvas.width, y);
|
|
ctx.stroke();
|
|
}
|
|
ctx.fillStyle = "#777";
|
|
ctx.font = "16px sans-serif";
|
|
ctx.textAlign = "center";
|
|
ctx.fillText("Load image/video background, then draw regions", this.canvas.width / 2, this.canvas.height / 2);
|
|
}
|
|
|
|
drawRegion(region, selected) {
|
|
const ctx = this.ctx;
|
|
const type = String(region.type || "rect").toLowerCase();
|
|
ctx.save();
|
|
ctx.strokeStyle = selected ? "#ffcc00" : "#35c759";
|
|
ctx.fillStyle = selected ? "rgba(255,204,0,0.20)" : "rgba(53,199,89,0.20)";
|
|
ctx.lineWidth = selected ? 3 : 2;
|
|
if (type === "rect") {
|
|
const rect = normalizeRect(region);
|
|
ctx.fillRect(rect.x, rect.y, rect.w, rect.h);
|
|
ctx.strokeRect(rect.x, rect.y, rect.w, rect.h);
|
|
} else if (type === "circle") {
|
|
ctx.beginPath();
|
|
ctx.arc(Number(region.cx ?? region.x ?? 0), Number(region.cy ?? region.y ?? 0), Number(region.r ?? region.radius ?? 0), 0, Math.PI * 2);
|
|
ctx.fill();
|
|
ctx.stroke();
|
|
} else if (type === "polygon") {
|
|
this.drawPolygonPath(region.points || []);
|
|
ctx.fill();
|
|
ctx.stroke();
|
|
}
|
|
ctx.restore();
|
|
}
|
|
|
|
drawPolygonPath(points) {
|
|
const ctx = this.ctx;
|
|
ctx.beginPath();
|
|
points.forEach((point, index) => {
|
|
const x = Array.isArray(point) ? point[0] : point.x;
|
|
const y = Array.isArray(point) ? point[1] : point.y;
|
|
if (index === 0) ctx.moveTo(x, y);
|
|
else ctx.lineTo(x, y);
|
|
});
|
|
ctx.closePath();
|
|
}
|
|
|
|
drawDraft(draft) {
|
|
const region = this.draftToRegion(draft);
|
|
if (!region) return;
|
|
this.drawRegion(region, true);
|
|
}
|
|
|
|
drawPolygonDraft() {
|
|
const ctx = this.ctx;
|
|
ctx.save();
|
|
ctx.strokeStyle = "#ffcc00";
|
|
ctx.fillStyle = "#ffcc00";
|
|
ctx.lineWidth = 2;
|
|
ctx.beginPath();
|
|
this.polygonDraft.forEach((point, index) => {
|
|
if (index === 0) ctx.moveTo(point.x, point.y);
|
|
else ctx.lineTo(point.x, point.y);
|
|
ctx.fillRect(point.x - 3, point.y - 3, 6, 6);
|
|
});
|
|
ctx.stroke();
|
|
ctx.restore();
|
|
}
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "ComfyUI.UniversalRemove.RegionEditor",
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData?.name !== "UR Region Editor") return;
|
|
chainCallback(nodeType.prototype, "onNodeCreated", function () {
|
|
this.serialize_widgets = true;
|
|
const editor = new RegionEditorWidget(this);
|
|
this.urRegionEditor = editor;
|
|
const widget = this.addDOMWidget("region_editor", "region_editor", editor.root);
|
|
widget.computeSize = (width) => [width, Math.max(420, this.size?.[1] ? this.size[1] - 160 : 420)];
|
|
this.setSize([Math.max(this.size?.[0] || 680, 680), Math.max(this.size?.[1] || 680, 680)]);
|
|
chainCallback(this, "onResize", function () {
|
|
editor.render();
|
|
});
|
|
});
|
|
},
|
|
});
|