Files
avatechai-avatar-graph-comfyui/js/LayerEditor.js
T

697 lines
21 KiB
JavaScript

import { SideBar } from "./SideBar.js";
import { api } from "./api.js";
import { app } from "./app.js";
import { runONNX } from "./onnx.js";
import {
showImageEditor,
point_label,
imageUrl,
imageContainerSize,
imagePrompts,
targetNode,
imageSize,
selectedLayer,
imagePromptsMulti,
embeddings,
embeddingID,
alertDialog,
allImagePrompts,
boxesMulti,
enableAutoSegment,
} from "./state.js";
import { van } from "./van.js";
import vision from "https://cdn.jsdelivr.net/npm/@mediapipe/tasks-vision@0.10.3";
const { PoseLandmarker, FaceLandmarker, FilesetResolver } = vision;
const { button, div, img, canvas, span } = van.tags;
let throttle = false;
const positivePrompt = van.state(true);
const enableBackgroundRemover = van.state(true);
const isMobileDevice = () => {
return window.screen.width < 768;
};
// Auto segmentation
const filesetResolver = await FilesetResolver.forVisionTasks(
"https://cdn.jsdelivr.net/npm/@mediapipe/tasks-vision@0.10.3/wasm"
);
const faceLandmarker = await FaceLandmarker.createFromOptions(filesetResolver, {
baseOptions: {
modelAssetPath: `https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task`,
delegate: "GPU",
},
// outputFaceBlendshapes: true,
runningMode: "IMAGE",
numFaces: 1,
});
const poseLandmarker = await PoseLandmarker.createFromOptions(filesetResolver, {
baseOptions: {
modelAssetPath: `https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/1/pose_landmarker_full.task`,
delegate: "GPU",
},
runningMode: "IMAGE",
numPoses: 1,
});
const layerMapping = {
L_eye: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: 0.25,
negativeScale: 0.5,
indices: FaceLandmarker.FACE_LANDMARKS_LEFT_EYE,
},
R_eye: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: 0.25,
negativeScale: 0.5,
indices: FaceLandmarker.FACE_LANDMARKS_RIGHT_EYE,
},
L_iris: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: -0.2,
negativeScale: 0.5,
indices: FaceLandmarker.FACE_LANDMARKS_LEFT_IRIS,
},
R_iris: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: -0.2,
negativeScale: 0.5,
indices: FaceLandmarker.FACE_LANDMARKS_RIGHT_IRIS,
},
face: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 60,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: 0.5,
negativeScale: 0,
indices: FaceLandmarker.FACE_LANDMARKS_FACE_OVAL,
},
mouth: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: -0.3,
negativeScale: 0.3,
// https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
indices: [61, 37, 270, 91, 314].map((x) => ({
start: x,
end: x,
})),
},
mouth_in: {
useMiddle: false,
positiveOffsetX: 0,
positiveOffsetY: 0,
negativeOffsetX: 0,
negativeOffsetY: 0,
positiveScale: -0.5,
negativeScale: 0.5,
// https://stackoverflow.com/questions/66649492/how-to-get-specific-landmark-of-face-like-lips-or-eyes-using-tensorflow-js-face
indices: [310, 88].map((x) => ({
start: x,
end: x,
})),
},
};
export const segmented = van.state(false);
export async function autoSegment() {
const image = document.getElementById("image");
const landmarks = faceLandmarker.detect(image).faceLandmarks[0];
Object.entries(layerMapping).forEach(([key, value]) => {
imagePromptsMulti.val[key] = [];
});
Object.entries(layerMapping).forEach(([key, value]) => {
const positivePoints = [];
const middlePoints = [];
const negativePoints = [];
// Positive points
for (const { start, end } of value.indices) {
const startPoint = landmarks[start];
// const endPoint = landmarks[end];
const startX = startPoint.x * imageSize.val.width;
const startY = startPoint.y * imageSize.val.height;
// const endX = endPoint.x * imageSize.val.width;
// const endY = endPoint.y * imageSize.val.height;
if (middlePoints.length === 0) {
middlePoints.push({ x: startX, y: startY, label: 1, isAuto: true });
// middlePoints.push({ x: endX, y: endY, label: 1 });
} else {
middlePoints[0].x += startX;
middlePoints[0].y += startY;
// middlePoints[1].x += endX;
// middlePoints[1].y += endY;
}
positivePoints.push({ x: startX, y: startY, label: 1, isAuto: true });
// positivePoints.push({ x: endX, y: endY, label: 1 });
// imagePrompts.val = [...imagePrompts.val, { x, y, label: 1 }];
}
// Middle points
const len = value.indices.length;
middlePoints[0].x /= len;
middlePoints[0].y /= len;
// middlePoints[1].x /= len;
// middlePoints[1].y /= len;
if (value.useMiddle) {
imagePromptsMulti.val[key] = [
...imagePromptsMulti.val[key],
...middlePoints,
];
} else {
// Negative points
for (const [i, { start, end }] of value.indices.entries()) {
const startPoint = landmarks[start];
// const endPoint = landmarks[end];
const startX = startPoint.x * imageSize.val.width;
const startY = startPoint.y * imageSize.val.height;
// const endX = endPoint.x * imageSize.val.width;
// const endY = endPoint.y * imageSize.val.height;
const middlePoint = middlePoints[0];
const directionVector = {
x: middlePoint.x - startX,
y: middlePoint.y - startY,
};
const directionVectorLength = Math.sqrt(
directionVector.x * directionVector.x +
directionVector.y * directionVector.y
);
if (value.negativeScale !== 0) {
const negativePointDistance =
value.negativeScale * directionVectorLength;
const negativePoint = {
x:
startX -
(negativePointDistance * directionVector.x) /
directionVectorLength -
value.negativeOffsetX,
y:
startY -
(negativePointDistance * directionVector.y) /
directionVectorLength -
value.negativeOffsetY,
label: 0,
isAuto: true,
};
negativePoints.push(negativePoint);
}
const positivePointDistance =
value.positiveScale * directionVectorLength;
positivePoints[i] = {
x:
positivePoints[i].x -
(positivePointDistance * directionVector.x) /
directionVectorLength -
value.positiveOffsetX,
y:
positivePoints[i].y -
(positivePointDistance * directionVector.y) /
directionVectorLength -
value.positiveOffsetY,
label: 1,
isAuto: true,
};
}
imagePromptsMulti.val[key] = [
...imagePromptsMulti.val[key],
...positivePoints,
...negativePoints,
];
}
// Find bounding box of positive/negative points
const points = negativePoints.length > 0 ? negativePoints : positivePoints;
const box = {
x1: Math.min(...points.map((x) => x.x)),
y1: Math.min(...points.map((x) => x.y)),
x2: Math.max(...points.map((x) => x.x)),
y2: Math.max(...points.map((x) => x.y)),
};
boxesMulti.val[key] = box;
});
const poseLandmarks = poseLandmarker.detect(image).landmarks[0];
const positiveBreathX =
((poseLandmarks[11].x + poseLandmarks[12].x) / 2) * imageSize.val.width;
const positiveBreathY =
((poseLandmarks[11].y + poseLandmarks[12].y) / 2) * imageSize.val.height;
const negativeBreathX1 = poseLandmarks[0].x * imageSize.val.width;
const negativeBreathY1 = poseLandmarks[0].y * imageSize.val.height;
const negativeBreathX2 = poseLandmarks[9].x * imageSize.val.width;
const negativeBreathY2 = poseLandmarks[9].y * imageSize.val.height;
const negativeBreathX3 = poseLandmarks[10].x * imageSize.val.width;
const negativeBreathY3 = poseLandmarks[10].y * imageSize.val.height;
imagePromptsMulti.val["breath"] = [
{ x: positiveBreathX, y: positiveBreathY, label: 1, isAuto: true },
{ x: negativeBreathX1, y: negativeBreathY1, label: 0, isAuto: true },
{ x: negativeBreathX2, y: negativeBreathY2, label: 0, isAuto: true },
{ x: negativeBreathX3, y: negativeBreathY3, label: 0, isAuto: true },
];
imagePrompts.val = imagePromptsMulti.val[selectedLayer.val];
segmented.val = true;
console.log("Done");
}
export function setRemoveBackgroundNode() {
const rmBgNodes = app.graph.findNodesByType(
"Image Rembg (Remove Background)"
);
if (!rmBgNodes?.length) {
alertDialog.val = {
text: "Remove background node not found. Please ensure the workflow is correct.",
time: 5000,
};
return;
}
rmBgNodes.forEach((node) => {
// node is bypassed if mode is 4
node.mode = enableBackgroundRemover.val ? 0 : 4;
});
}
export function updateImagePrompts() {
if (selectedLayer.val !== "" && selectedLayer.val !== undefined) {
imagePromptsMulti.val = {
...imagePromptsMulti.val,
[selectedLayer.val]: imagePrompts.val,
};
targetNode.val.widgets.find((x) => x.name === "image_prompts_json").value =
JSON.stringify(imagePromptsMulti.val);
// const canvas = document.getElementById("mask-canvas");
// const base64Image = canvas.toDataURL();
// api.fetchApi("/segments", {
// method: "POST",
// body: JSON.stringify({
// name: embeddingID.val,
// segments: {
// [selectedLayer.val]: base64Image,
// },
// }),
// });
} else {
targetNode.val.widgets.find((x) => x.name === "image_prompts_json").value =
JSON.stringify(imagePrompts.val);
}
targetNode.val.graph.change();
}
export async function uploadSegments() {
const emptyLayers = [];
Object.entries(imagePromptsMulti.val).forEach(([key, value]) => {
if (value.length === 0) {
emptyLayers.push(key);
}
});
if (emptyLayers.length > 0) {
alertDialog.val = {
text: "The following layers have no segments: " + emptyLayers.join(", "),
time: 5000,
};
return false;
}
const segments = {};
for (const [layer, prompts] of Object.entries(imagePromptsMulti.val)) {
await drawSegment(getClicks(prompts), layer, false);
const canvas = document.getElementById("mask-canvas");
const base64Image = canvas.toDataURL();
segments[layer] = base64Image;
// download image
// const a = document.createElement("a");
// a.href = base64Image;
// a.download = layer + ".png";
// a.click();
}
await api.fetchApi("/segments", {
method: "POST",
body: JSON.stringify({
name: embeddingID.val,
segments,
}),
});
return true;
}
async function handleClick(e) {
const rect = e.target.getBoundingClientRect();
const x = e.clientX - rect.left;
const y = e.clientY - rect.top;
const relativeX = Math.trunc(
((x / e.target.offsetWidth) * imageSize.val.width) / imageSize.val.imgScale
);
const relativeY = Math.trunc(
((y / e.target.offsetHeight) * imageSize.val.height) /
imageSize.val.imgScale
);
let label;
if (isMobileDevice()) {
label = positivePrompt.val ? 1 : 0;
} else {
label = e.isRight ? 0 : 1;
}
imagePrompts.val = [
...imagePrompts.val,
{ x: relativeX, y: relativeY, label },
];
await drawSegment(getClicks());
updateImagePrompts();
}
async function handlePointClick(e, point) {
e.preventDefault();
imagePrompts.val = imagePrompts.val.filter(
(x) => !(x.x === point.x && x.y === point.y)
);
await drawSegment(getClicks());
updateImagePrompts();
}
function handleImageSize(image) {
// Input images to SAM must be resized so the longest side is 1024
const documentHeight = document.documentElement.clientHeight;
const LONG_SIDE_LENGTH = 1024;
let w = image.naturalWidth;
let h = image.naturalHeight;
const samScale = LONG_SIDE_LENGTH / Math.max(h, w);
const imgScale = documentHeight / Math.max(h, w);
return { height: h, width: w, samScale, imgScale };
}
export function getClicks(prompts) {
return (prompts || imagePrompts.val).map((point) => ({
x: point.x,
y: point.y,
clickType: point.label,
isAuto: point.isAuto,
}));
}
export async function drawSegment(clicks, layer, drawBox = true) {
const canvas = document.getElementById("mask-canvas");
const ctx = canvas.getContext("2d");
if (clicks.length === 0) {
ctx.clearRect(0, 0, canvas.width, canvas.height);
return;
}
if (embeddings.val) {
const box = enableAutoSegment.val
? boxesMulti.val[layer || selectedLayer.val]
: null;
const filteredClicks = enableAutoSegment.val
? clicks
: clicks.filter((click) => !click.isAuto);
if (filteredClicks.length === 0) {
ctx.clearRect(0, 0, canvas.width, canvas.height);
return;
}
const mask = await runONNX(filteredClicks, embeddings.val, box);
if (mask) {
ctx.clearRect(0, 0, canvas.width, canvas.height);
ctx.drawImage(mask, 0, 0);
if (box && drawBox) {
ctx.strokeStyle = "green";
ctx.lineWidth = 5;
ctx.strokeRect(box.x1, box.y1, box.x2 - box.x1, box.y2 - box.y1);
}
}
}
}
export function LayerEditor() {
let realTimeSegment = true;
const showSidebar = van.state(true);
document.addEventListener("keydown", (e) => {
if (showImageEditor.val && e.code === "Tab") {
e.preventDefault();
realTimeSegment = !realTimeSegment;
if (!realTimeSegment) {
drawSegment(getClicks());
}
}
});
return div(
{
class: () =>
"absolute flex bg-gray-900 bg-opacity-50 top-0 w-full h-full pointer-events-auto z-[1000] " +
(showImageEditor.val ? "" : "hidden"),
},
div(
{
class:
"absolute top-4 left-4 right-0 flex w-full gap-2 justify-start z-[200]",
},
button(
{
class: () => "btn btn-neutral flex flex-row normal-case rounded-md",
onclick: async () => {
console.log("close");
showImageEditor.val = false;
await uploadSegments();
const isEqual = allImagePrompts.val.map(
(x) =>
JSON.stringify(imagePromptsMulti.val) ===
JSON.stringify(x.prompt)
);
if (!isEqual.includes(true))
allImagePrompts.val = [
...allImagePrompts.val,
{
version: "v" + allImagePrompts.val.length,
prompt: imagePromptsMulti.val,
},
];
// api.fetchApi("/segments_order", {
// method: "POST",
// body: JSON.stringify({
// name: embeddingID.val,
// order: Object.keys(imagePromptsMulti.val),
// }),
// });
},
},
span({
class: "iconify text-lg",
"data-icon": "ic:baseline-arrow-back",
"data-inline": "false",
}),
div("Back")
),
button(
{
class: () => "btn btn-neutral flex flex-row normal-case rounded-md",
onclick: () => (showSidebar.val = !showSidebar.val),
},
div(() => (showSidebar.val ? "Hide UI" : "Show UI"))
),
button(
{
class: () => "btn btn-neutral flex flex-row normal-case rounded-md",
onclick: () => {
enableAutoSegment.val = !enableAutoSegment.val;
drawSegment(getClicks());
},
},
() => (enableAutoSegment.val ? "Auto Segment On" : "Auto Segment Off")
),
button(
{
class: () => "btn btn-neutral flex flex-row normal-case rounded-md",
onclick: () => {
enableBackgroundRemover.val = !enableBackgroundRemover.val;
setRemoveBackgroundNode();
},
},
() =>
enableBackgroundRemover.val
? "Background Remover On"
: "Background Remover Off"
),
button(
{
class: () =>
`btn btn-neutral flex flex-row normal-case rounded-md ${
isMobileDevice() ? "" : "hidden"
}`,
onclick: () => (positivePrompt.val = !positivePrompt.val),
},
div(() => (positivePrompt.val ? "Positive" : "Negative"))
)
),
div(
{
class:
"hidden w-full justify-center absolute top-0 left-0 right-0 items-center",
},
button(
{
class: () => " px-4 py-2 rounded-md left-0 top-0 z-[200] ",
onclick: () => {
point_label.val = 1;
},
},
"Positive"
),
button(
{
class: () => " px-4 py-2 rounded-md left-0 top-0 z-[200] ",
onclick: () => {
point_label.val = 0;
},
},
"Negative"
)
),
div(
{
class: "flex items-center justify-center w-full h-full",
id: "image-container",
},
img({
id: "image",
class:
"fixed top-1/2 left-1/2 transform -translate-x-1/2 -translate-y-1/2",
src: imageUrl,
onload: async (e) => {
imageSize.val = handleImageSize(e.target);
document.getElementById("image-container").style.scale =
imageSize.val.imgScale;
imageContainerSize.val = {
width: e.target.offsetWidth,
height: e.target.offsetHeight,
};
const canvas = document.getElementById("mask-canvas");
canvas.width = e.target.naturalWidth;
canvas.height = e.target.naturalHeight;
},
oncontextmenu: async (e) => {
e.preventDefault();
e.isRight = true;
await handleClick(e);
},
onclick: async (e) => {
await handleClick(e);
},
onmouseleave: (e) => {
drawSegment(getClicks());
},
onmousemove: (e) => {
if (!throttle) {
throttle = true;
setTimeout(() => {
throttle = false;
if (embeddings.val && realTimeSegment) {
const rect = e.target.getBoundingClientRect();
const x = e.clientX - rect.left;
const y = e.clientY - rect.top;
const relativeX = Math.trunc(
((x / e.target.offsetWidth) * imageSize.val.width) /
imageSize.val.imgScale
);
const relativeY = Math.trunc(
((y / e.target.offsetHeight) * imageSize.val.height) /
imageSize.val.imgScale
);
const clicks = [
...getClicks(),
{ x: relativeX, y: relativeY, clickType: 1 },
];
drawSegment(clicks);
}
}, 10);
}
},
}),
() =>
canvas({
class:
"pointer-events-none fixed top-1/2 left-1/2 transform -translate-x-1/2 -translate-y-1/2 opacity-80",
style: () =>
`width: ${imageContainerSize.val.width}px; height: ${imageContainerSize.val.height}px;`,
id: "mask-canvas",
}),
() =>
div(
{
class: "absolute w-full h-full pointer-events-none",
style: () =>
`width: ${imageContainerSize.val.width}px; height: ${imageContainerSize.val.height}px;`,
},
...(enableAutoSegment.val
? imagePrompts.val
: imagePrompts.val?.filter((click) => !click.isAuto)
).map((point) => {
return button({
style: () =>
`left: ${
(point.x / imageSize.val.width) * imageContainerSize.val.width
}px; top: ${
(point.y / imageSize.val.height) *
imageContainerSize.val.height
}px; transform: translate(-50%, -50%);`,
class: () =>
`absolute w-3 h-3 p-0 rounded-full pointer-events-auto ${
point.label === 1 ? "bg-green-500" : "bg-red-500"
}`,
oncontextmenu: async (e) => {
await handlePointClick(e, point);
},
onclick: async (e) => {
await handlePointClick(e, point);
},
});
})
)
),
() => (showSidebar.val ? SideBar() : div())
);
}