351 lines
11 KiB
JavaScript
351 lines
11 KiB
JavaScript
import * as THREE from "three";
|
|
import { OrbitControls } from "../vendor/OrbitControls.mjs";
|
|
import { GLTFExporter } from "../vendor/GLTFExporter.mjs";
|
|
import { OBJExporter } from "../vendor/OBJExporter.mjs";
|
|
|
|
const SOURCE = "gokayfem.depth-visualization";
|
|
const container = document.querySelector("#canvas-container");
|
|
const statusElement = document.querySelector("#status");
|
|
const errorElement = document.querySelector("#error");
|
|
const batchSelect = document.querySelector("#batch-select");
|
|
const depthScale = document.querySelector("#depth-scale");
|
|
const depthValue = document.querySelector("#depth-value");
|
|
const exportButton = document.querySelector("#export-mesh");
|
|
|
|
const renderer = new THREE.WebGLRenderer({
|
|
antialias: true,
|
|
alpha: false,
|
|
preserveDrawingBuffer: true,
|
|
powerPreference: "high-performance",
|
|
});
|
|
renderer.setPixelRatio(Math.min(window.devicePixelRatio || 1, 2));
|
|
renderer.outputColorSpace = THREE.SRGBColorSpace;
|
|
container.append(renderer.domElement);
|
|
|
|
const scene = new THREE.Scene();
|
|
scene.background = new THREE.Color(0x111318);
|
|
|
|
const camera = new THREE.PerspectiveCamera(42, 1, 0.01, 1000);
|
|
const controls = new OrbitControls(camera, renderer.domElement);
|
|
controls.enableDamping = true;
|
|
controls.dampingFactor = 0.08;
|
|
controls.screenSpacePanning = true;
|
|
|
|
scene.add(new THREE.HemisphereLight(0xffffff, 0x303040, 2.2));
|
|
const keyLight = new THREE.DirectionalLight(0xffffff, 1.5);
|
|
keyLight.position.set(4, 6, 8);
|
|
scene.add(keyLight);
|
|
|
|
let channel = null;
|
|
let viewUrl = null;
|
|
let output = null;
|
|
let mesh = null;
|
|
let updateVersion = 0;
|
|
let animationFrame = null;
|
|
let disposed = false;
|
|
|
|
function setStatus(message) {
|
|
statusElement.textContent = message;
|
|
statusElement.hidden = false;
|
|
errorElement.hidden = true;
|
|
}
|
|
|
|
function setError(error) {
|
|
console.error("[Depth Viewer]", error);
|
|
errorElement.textContent = error instanceof Error ? error.message : String(error);
|
|
errorElement.hidden = false;
|
|
statusElement.hidden = true;
|
|
}
|
|
|
|
function resetCamera() {
|
|
camera.position.set(0, 0, 9);
|
|
controls.target.set(0, 0, 0);
|
|
controls.update();
|
|
}
|
|
|
|
function resize() {
|
|
const width = Math.max(container.clientWidth, 1);
|
|
const height = Math.max(container.clientHeight, 1);
|
|
renderer.setSize(width, height, false);
|
|
camera.aspect = width / height;
|
|
camera.updateProjectionMatrix();
|
|
}
|
|
|
|
new ResizeObserver(resize).observe(container);
|
|
resetCamera();
|
|
resize();
|
|
|
|
function disposeMaterial(material) {
|
|
for (const value of Object.values(material)) {
|
|
if (value?.isTexture) {
|
|
value.dispose();
|
|
}
|
|
}
|
|
material.dispose();
|
|
}
|
|
|
|
function removeMesh() {
|
|
if (!mesh) {
|
|
return;
|
|
}
|
|
scene.remove(mesh);
|
|
mesh.geometry.dispose();
|
|
disposeMaterial(mesh.material);
|
|
mesh = null;
|
|
exportButton.disabled = true;
|
|
}
|
|
|
|
function imageUrl(descriptor) {
|
|
if (!viewUrl) {
|
|
throw new Error("The ComfyUI API URL has not been initialized.");
|
|
}
|
|
const url = new URL(viewUrl, window.location.origin);
|
|
url.search = new URLSearchParams({
|
|
filename: descriptor.filename,
|
|
subfolder: descriptor.subfolder ?? "",
|
|
type: descriptor.type ?? "temp",
|
|
}).toString();
|
|
return url.href;
|
|
}
|
|
|
|
function loadTexture(descriptor, colorTexture) {
|
|
return new Promise((resolve, reject) => {
|
|
new THREE.TextureLoader().load(
|
|
imageUrl(descriptor),
|
|
(texture) => {
|
|
texture.colorSpace = colorTexture
|
|
? THREE.SRGBColorSpace
|
|
: THREE.NoColorSpace;
|
|
texture.anisotropy = Math.min(
|
|
8,
|
|
renderer.capabilities.getMaxAnisotropy(),
|
|
);
|
|
resolve(texture);
|
|
},
|
|
undefined,
|
|
() => reject(new Error(`Unable to load ${descriptor.filename}.`)),
|
|
);
|
|
});
|
|
}
|
|
|
|
async function showFrame(index) {
|
|
if (!output) {
|
|
return;
|
|
}
|
|
const reference = output.reference_image[index] ?? output.reference_image[0];
|
|
const depth = output.depth_map[index] ?? output.depth_map[0];
|
|
if (!reference || !depth) {
|
|
setError("The selected frame is missing a reference image or depth map.");
|
|
return;
|
|
}
|
|
|
|
const version = ++updateVersion;
|
|
setStatus(`Loading frame ${index + 1}…`);
|
|
try {
|
|
const [referenceTexture, depthTexture] = await Promise.all([
|
|
loadTexture(reference, true),
|
|
loadTexture(depth, false),
|
|
]);
|
|
if (version !== updateVersion || disposed) {
|
|
referenceTexture.dispose();
|
|
depthTexture.dispose();
|
|
return;
|
|
}
|
|
|
|
removeMesh();
|
|
const image = referenceTexture.image;
|
|
const aspect = image.width / Math.max(image.height, 1);
|
|
const width = 7;
|
|
const height = width / aspect;
|
|
const geometry = new THREE.PlaneGeometry(width, height, 256, 256);
|
|
const material = new THREE.MeshStandardMaterial({
|
|
map: referenceTexture,
|
|
displacementMap: depthTexture,
|
|
displacementScale: Number(depthScale.value),
|
|
displacementBias: -Number(depthScale.value) / 2,
|
|
roughness: 0.95,
|
|
metalness: 0,
|
|
side: THREE.DoubleSide,
|
|
});
|
|
mesh = new THREE.Mesh(geometry, material);
|
|
scene.add(mesh);
|
|
resetCamera();
|
|
statusElement.hidden = true;
|
|
exportButton.disabled = false;
|
|
} catch (error) {
|
|
if (version === updateVersion) {
|
|
removeMesh();
|
|
setError(error);
|
|
}
|
|
}
|
|
}
|
|
|
|
function setOutput(nextOutput) {
|
|
const referenceCount = nextOutput?.reference_image?.length ?? 0;
|
|
const depthCount = nextOutput?.depth_map?.length ?? 0;
|
|
const count = Math.max(referenceCount, depthCount);
|
|
if (!count) {
|
|
setError("ComfyUI returned no depth-viewer images.");
|
|
return;
|
|
}
|
|
|
|
output = nextOutput;
|
|
batchSelect.replaceChildren();
|
|
for (let index = 0; index < count; index += 1) {
|
|
const option = document.createElement("option");
|
|
option.value = String(index);
|
|
option.textContent = `${index + 1} / ${count}`;
|
|
batchSelect.append(option);
|
|
}
|
|
batchSelect.disabled = count === 1;
|
|
batchSelect.value = "0";
|
|
void showFrame(0);
|
|
}
|
|
|
|
function download(blob, filename) {
|
|
const url = URL.createObjectURL(blob);
|
|
const link = document.createElement("a");
|
|
link.href = url;
|
|
link.download = filename;
|
|
link.click();
|
|
setTimeout(() => URL.revokeObjectURL(url), 1000);
|
|
}
|
|
|
|
function bakedMesh() {
|
|
if (!mesh) {
|
|
throw new Error("Queue the node before exporting a mesh.");
|
|
}
|
|
const clone = new THREE.Mesh(
|
|
mesh.geometry.clone(),
|
|
new THREE.MeshStandardMaterial({
|
|
map: mesh.material.map,
|
|
roughness: mesh.material.roughness,
|
|
metalness: mesh.material.metalness,
|
|
side: THREE.DoubleSide,
|
|
}),
|
|
);
|
|
|
|
const depthImage = mesh.material.displacementMap.image;
|
|
const canvas = document.createElement("canvas");
|
|
canvas.width = depthImage.width;
|
|
canvas.height = depthImage.height;
|
|
const context = canvas.getContext("2d", { willReadFrequently: true });
|
|
context.drawImage(depthImage, 0, 0);
|
|
const pixels = context.getImageData(0, 0, canvas.width, canvas.height).data;
|
|
const positions = clone.geometry.attributes.position;
|
|
const uvs = clone.geometry.attributes.uv;
|
|
const scale = mesh.material.displacementScale;
|
|
const bias = mesh.material.displacementBias;
|
|
|
|
for (let index = 0; index < positions.count; index += 1) {
|
|
const x = Math.min(
|
|
canvas.width - 1,
|
|
Math.max(0, Math.round(uvs.getX(index) * (canvas.width - 1))),
|
|
);
|
|
const y = Math.min(
|
|
canvas.height - 1,
|
|
Math.max(0, Math.round((1 - uvs.getY(index)) * (canvas.height - 1))),
|
|
);
|
|
const depth = pixels[(y * canvas.width + x) * 4] / 255;
|
|
positions.setZ(index, depth * scale + bias);
|
|
}
|
|
positions.needsUpdate = true;
|
|
clone.geometry.computeVertexNormals();
|
|
return clone;
|
|
}
|
|
|
|
async function exportMesh() {
|
|
const format = document.querySelector("#export-format").value;
|
|
const exportable = bakedMesh();
|
|
try {
|
|
if (format === "obj") {
|
|
const data = new OBJExporter().parse(exportable);
|
|
download(new Blob([data], { type: "text/plain" }), "depth-mesh.obj");
|
|
return;
|
|
}
|
|
|
|
const binary = format === "glb";
|
|
const data = await new GLTFExporter().parseAsync(exportable, { binary });
|
|
const blob = binary
|
|
? new Blob([data], { type: "model/gltf-binary" })
|
|
: new Blob([JSON.stringify(data, null, 2)], {
|
|
type: "model/gltf+json",
|
|
});
|
|
download(blob, `depth-mesh.${format}`);
|
|
} finally {
|
|
exportable.geometry.dispose();
|
|
exportable.material.dispose();
|
|
}
|
|
}
|
|
|
|
function takeScreenshot() {
|
|
renderer.render(scene, camera);
|
|
renderer.domElement.toBlob((blob) => {
|
|
if (blob) {
|
|
download(blob, "depth-preview.png");
|
|
}
|
|
}, "image/png");
|
|
}
|
|
|
|
function animate() {
|
|
if (disposed) {
|
|
return;
|
|
}
|
|
animationFrame = requestAnimationFrame(animate);
|
|
if (document.visibilityState === "visible") {
|
|
controls.update();
|
|
renderer.render(scene, camera);
|
|
}
|
|
}
|
|
animate();
|
|
|
|
batchSelect.addEventListener("change", () => {
|
|
void showFrame(Number(batchSelect.value));
|
|
});
|
|
depthScale.addEventListener("input", () => {
|
|
const value = Number(depthScale.value);
|
|
depthValue.value = value.toFixed(2);
|
|
if (mesh) {
|
|
mesh.material.displacementScale = value;
|
|
mesh.material.displacementBias = -value / 2;
|
|
}
|
|
});
|
|
document.querySelector("#reset-camera").addEventListener("click", resetCamera);
|
|
document.querySelector("#screenshot").addEventListener("click", takeScreenshot);
|
|
exportButton.addEventListener("click", () => {
|
|
void exportMesh().catch(setError);
|
|
});
|
|
|
|
window.addEventListener("message", (event) => {
|
|
if (
|
|
event.origin !== window.location.origin
|
|
|| event.source !== window.parent
|
|
|| event.data?.source !== SOURCE
|
|
) {
|
|
return;
|
|
}
|
|
if (event.data.type === "connect") {
|
|
channel = event.data.channel;
|
|
window.parent.postMessage(
|
|
{ source: SOURCE, channel, type: "ready" },
|
|
window.location.origin,
|
|
);
|
|
return;
|
|
}
|
|
if (event.data.channel !== channel) {
|
|
return;
|
|
}
|
|
if (event.data.type === "initialize") {
|
|
viewUrl = event.data.viewUrl;
|
|
} else if (event.data.type === "update") {
|
|
setOutput(event.data.output);
|
|
} else if (event.data.type === "dispose") {
|
|
disposed = true;
|
|
updateVersion += 1;
|
|
cancelAnimationFrame(animationFrame);
|
|
removeMesh();
|
|
controls.dispose();
|
|
renderer.dispose();
|
|
}
|
|
});
|