Public
Edited
Nov 26, 2024
Fork of Three.js
1 star
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
// 2D view
{
const {
imageMesh,
channelMesh,
outputMesh,
convs,
sums,
min,
max,
outputBuf
} = outputs
const offscreen = new OffscreenCanvas(32, 32)
const octx = offscreen.getContext('2d')

const canvas = document.createElement('canvas')
canvas.width = width
canvas.height = 256
const ctx = canvas.getContext('2d', { willReadFrequently: true })
const outputIm = ctx.createImageData(32, 32)
for (let i = 0; i < outputBuf.length; i += 4) {
for (let j = 0; j < 4; j++) {
outputIm.data[i + j] = outputBuf[i + j] * 255
}
}
octx.putImageData(outputIm, 0, 0)
// clean up old connectors
const clearConnectors = () => {
for (let i = tubeGroup.children.length - 1; i >= 0; i--) {
const child = tubeGroup.children[i]
child.geometry.dispose()
child.material.dispose()
tubeGroup.remove(child)
}
for (let i = sumGroup.children.length - 1; i >= 0; i--) {
const child = sumGroup.children[i]
child.geometry.dispose()
child.material.dispose()
sumGroup.remove(child)
}
}

const drawImages = (kx, ky) => {
ctx.clearRect(0, 0, 256 * 2, 256)

ctx.lineWidth = 1
ctx.drawImage(image, 0, 0, 256, 256)
ctx.strokeStyle = `rgba(0, 0, 0, 0.5)`
for (let x = 0; x < image.width; x++) {
ctx.beginPath()
ctx.moveTo(x * 256 / 32, 0)
ctx.lineTo(x * 256 / 32, 256)
ctx.stroke()
}
for (let y = 0; y < image.height; y++) {
ctx.beginPath()
ctx.moveTo(0, y * 256 / 32)
ctx.lineTo(256, y * 256 / 32)
ctx.stroke()
}
ctx.drawImage(offscreen, 256, 0, 256, 256)
for (let x = 0; x < image.width; x++) {
ctx.beginPath()
ctx.moveTo(x * 256 / 32 + 256, 0)
ctx.lineTo(x * 256 / 32 + 256, 256)
ctx.stroke()
}
for (let y = 0; y < image.height; y++) {
ctx.beginPath()
ctx.moveTo(0 + 256, y * 256 / 32)
ctx.lineTo(256 + 256, y * 256 / 32)
ctx.stroke()
}
}

drawImages()

// reusable
const matrix = new THREE.Matrix4()
const vec3 = new THREE.Vector3()
const color = new THREE.Color()
// for matrix decomposition
const pos = new THREE.Vector3()
const quat = new THREE.Quaternion()
const scale = new THREE.Vector3()

let prev = {
index: null,
inputPxColor: new THREE.Color(),
outputPxColor: new THREE.Color(),
}
ctx.canvas.addEventListener('pointermove', e => {
const rect = ctx.canvas.getBoundingClientRect()
const x = Math.floor((e.clientX - rect.left) / 256 * 32)
const y = Math.floor((e.clientY - rect.top) / 256 * 32)
if (x >= 32 || y >= 32) {
return;
}
const index = x + y * 32;

if (index === prev.index) {
return;
}

clearConnectors()

if (prev.index !== null) {
imageMesh.setColorAt(prev.index, prev.inputPxColor)
outputMesh.setColorAt(prev.index, prev.outputPxColor)
}
prev.index = index
// only works when input image is same size as output image
imageMesh.getColorAt(index, prev.inputPxColor)
outputMesh.getColorAt(index, prev.outputPxColor)
const imagePos = new THREE.Vector3()
imageMesh.getMatrixAt(index, matrix)
matrix.decompose(imagePos, quat, scale)
imagePos.add(imageGroup.position)

const outputPos = new THREE.Vector3()
outputMesh.getMatrixAt(index, matrix)
matrix.decompose(outputPos, quat, scale)
outputPos.add(outputGroup.position)

// we could definitely optimize the creation of geoms/mats below,
// but this is performant enough for a demo
const sumPositions = []
const sumGeom = new THREE.BoxGeometry(kernel_dims[0], kernel_dims[1], boxSize)
const rgb = outputBuf.slice(index * 4, index * 4 + 3)
rgb.forEach((c, i) => {
const pos = new THREE.Vector3(
util.lerp(0, 3 - 1, i, ...x_range) + padded_image_dims[1] / 2,
padded_image_dims[1] / 2,
15
).add(depthGroup.position)
const rgb = [0, 0, 0]
rgb[i] = c
const sumMat = new THREE.MeshBasicMaterial({ color: color.setRGB(...rgb) })
const sum = new THREE.Mesh(sumGeom, sumMat)
sum.position.copy(pos)
sumPositions.push(sum.position.clone())
sumGroup.add(sum)

const pts = [
pos,
new THREE.Vector3(pos.x, pos.y, pos.z + 1),
util.pointAt(pos, outputPos, 0.5),
new THREE.Vector3(outputPos.x, outputPos.y, outputPos.z - 1),
outputPos
]
const curve = new THREE.CatmullRomCurve3(pts, false, 'centripetal', 0.5)
const geom = new THREE.TubeBufferGeometry(curve, 48, 0.2, 8, false)
const mat = new THREE.MeshBasicMaterial({ color })
tubeGroup.add(new THREE.Mesh(geom, mat))
})
let kernelIndices = []
for (let wx = x; wx < x + kernel_dims[0]; wx++) {
for (let wy = y; wy < y + kernel_dims[1]; wy++) {
const i = (wx * 3) + (wy * 3 * padded_image_dims[0])
kernelIndices.push(i + 0, i + 1, i + 2)
}
}

for (let i = 0; i < kernelIndices.length; i++) {
channelMesh.getColorAt(kernelIndices[i], color)
channelMesh.getMatrixAt(kernelIndices[i], matrix)
matrix.decompose(pos, quat, scale)
pos.add(depthGroup.position)
const sumPos = sumPositions[i % 3]

// could just create these once and cache them, since
// they have the same curves/positions, just offset...
{
const pts = [
imagePos,
new THREE.Vector3(imagePos.x, imagePos.y, imagePos.z + 1),
util.pointAt(imagePos, pos, 0.5),
new THREE.Vector3(pos.x, pos.y, pos.z - 2),
pos, // end
]
const curve = new THREE.CatmullRomCurve3(pts, false, 'centripetal', 0.5)
const geom = new THREE.TubeBufferGeometry(curve, 48, 0.05, 8, false)
const mat = new THREE.MeshBasicMaterial({ color })
tubeGroup.add(new THREE.Mesh(geom, mat))
}

{
const pts = [
pos,
new THREE.Vector3(pos.x, pos.y, pos.z + 1),
util.pointAt(pos, sumPos, 0.5),
new THREE.Vector3(sumPos.x, sumPos.y, sumPos.z - 1),
sumPos
]
const curve = new THREE.CatmullRomCurve3(pts, false, 'centripetal', 0.5)
const geom = new THREE.TubeBufferGeometry(curve, 48, 0.05, 8, false)
const mat = new THREE.MeshBasicMaterial({ color })
tubeGroup.add(new THREE.Mesh(geom, mat))
}
}

imageMesh.setColorAt(index, color.setRGB(1, 1, 1))
imageMesh.instanceColor.needsUpdate = true
outputMesh.setColorAt(index, color)
outputMesh.instanceColor.needsUpdate = true
channelMesh.instanceColor.needsUpdate = true;
drawImages(x, y)
const d = ctx.getImageData(x / 32 * 256, y / 32 * 256, 1, 1).data;
ctx.fillStyle = `rgb(${d[0]}, ${d[1]}, ${d[2]})`
ctx.fillRect(512, 0, 256, 256)

ctx.fillStyle = '#fff'
ctx.fillRect(x / 32 * 256, y / 32 * 256, 256 / 32, 256 / 32)
})

ctx.canvas.addEventListener('pointerout', () => {
clearConnectors()
drawImages()
imageMesh.setColorAt(prev.index, prev.inputPxColor)
imageMesh.instanceColor.needsUpdate = true
outputMesh.setColorAt(prev.index, prev.outputPxColor)
outputMesh.instanceColor.needsUpdate = true
})
return ctx.canvas
}
Insert cell
Insert cell
Insert cell
pad = (kernel_dims[0] - 1) / 2
Insert cell
padded_image_dims = [image.width + pad * 2, image.height + pad * 2, 3]
Insert cell
output_image_dims = [
Math.floor((padded_image_dims[0] - kernel_dims[0]) / stride + 1),
Math.floor((padded_image_dims[1] - kernel_dims[1]) / stride + 1),
padded_image_dims[2] // rgb
]
Insert cell
Insert cell
outputs = {
// cleanup when cell re-runs
;[imageGroup, depthGroup, outputGroup].forEach(g => {
for (let i = g.children.length - 1; i >= 0; i--) {
const child = g.children[i]
child.geometry.dispose()
child.material.dispose()
g.remove(child)
}
})
imageGroup.position.set(-image.width / 2, -image.height / 2, -50)
depthGroup.position.set(-padded_image_dims[0] / 2, -padded_image_dims[1] / 2, -20)
outputGroup.position.set(-output_image_dims[0] / 2, -output_image_dims[1] / 2, 20)

// reusable
const matrix = new THREE.Matrix4()
const vec3 = new THREE.Vector3()
const color = new THREE.Color()
const geom = new THREE.BoxGeometry(boxSize, boxSize, boxSize)
const mat = new THREE.MeshBasicMaterial()

// original image
const imageMesh = new THREE.InstancedMesh(geom, mat, image.width * image.height)
imageGroup.add(imageMesh)

for (let i = 0, len = original_image_data.length; i < len; i += 4) {
const index = i / 4
const [x, y] = util.arrXY(index, image.width)
matrix.setPosition(vec3.set(x, image.height - y, 0))
imageMesh.setMatrixAt(index, matrix)
color.setRGB(
original_image_data[i + 0] / 255,
original_image_data[i + 1] / 255,
original_image_data[i + 2] / 255
)
imageMesh.setColorAt(index, color)
}
// depth / channels / rgb
const channelMesh = new THREE.InstancedMesh(geom, mat, padded_image_dims[0] * padded_image_dims[1] * 3)
depthGroup.add(channelMesh)
for (let i = 0, count = 0, len = padded_image_data.length; i < len; i += 4) {
const index = i / 4
const [x, y] = util.arrXY(index, padded_image_dims[0])
for (let j = 0; j < 3; j++) {
matrix.setPosition(
vec3.set(
x + util.lerp(0, 3 - 1, j, ...x_range),
padded_image_dims[1] - y,
0
)
)
channelMesh.setMatrixAt(count, matrix)
const rgb = [0, 0, 0]
rgb[j] = padded_image_data[i + j] / 255
color.setRGB(...rgb)
channelMesh.setColorAt(count, color)
count++
}
}

// output
const outputMesh = new THREE.InstancedMesh(geom, mat, output_image_dims[0] * output_image_dims[1])
outputGroup.add(outputMesh)
// window

const convs = [] // w * k
const sums = [] // one r, g, or b value
const out = []
// note that for performance reasons, we would normally wouldn't use
// nested for-loops and would only calculate on indices, but this
// notebook is meant to be illustrative

// for each output pixel
for (let i = 0; i < output_image_dims[0] * output_image_dims[1]; i++) {
const [x, y] = util.arrXY(i, output_image_dims[0])

// calc window pos
const wx = x * stride
const wy = y * stride

let w = [[], [], []] // r,g,b channels; size 3x(5*5) = r,g,b x window
for (let j = 0; j < kernel_dims[0] * kernel_dims[1]; j++) {
const [kx, ky] = util.arrXY(j, kernel_dims[0])
const i = (wx + kx) + (wy + ky) * padded_image_dims[0]
// z = depth = rgb channel
for (let z = 0; z < 3; z++) {
w[z].push(padded_image_data[i * 4 + z])
}
}
const conv = w.map(depth => util.hadamard(depth, kernel)) // window * kernel
convs.push(conv)

// sum(wk) gives one value for one rgb component; norm this to get output image
sums.push(conv.map(kw => kw.reduce((sum, c) => sum + c, 0)))
}

let min = undefined
let max = undefined
for (let c of sums.flat()) {
if (min === undefined || c < min) min = c;
if (max === undefined || c > max) max = c;
}
for (let i = 0; i < sums.length; i++) {
const [wx, wy] = util.arrXY(i, output_image_dims[0])
for (let j = 0; j < kernel_dims[0] * kernel_dims[1]; j++) {
const [jx, jy] = util.arrXY(j, kernel_dims[0])
const px = wx + jx
const py = wy + jy
const oi = px + py * output_image_dims[1]
matrix.setPosition(vec3.set(px, output_image_dims[1] - py, 0))
outputMesh.setMatrixAt(oi, matrix)
color.setRGB(...sums[i].map(c => util.lerp(min, max, c)))
outputMesh.setColorAt(oi, color)
}
}
return {
imageMesh,
channelMesh,
outputMesh,
convs,
sums,
min,
max,
outputBuf: sums.flatMap(px => [...px.map(c => util.lerp(min, max, c)), 1]) // rgba
}
}
Insert cell
original_image_data = {
const canvas = new OffscreenCanvas(image.width, image.height)
const ctx = canvas.getContext('2d')
ctx.drawImage(image, 0, 0)
return ctx.getImageData(0, 0, image.width, image.height).data
}
Insert cell
padded_image_data = {
const canvas = new OffscreenCanvas(padded_image_dims[0], padded_image_dims[1])
const ctx = canvas.getContext('2d')
// ctx.scale(1, -1)
// ctx.drawImage(image, 0, 0, image.width, image.height, pad, -pad, image.width, -image.height)
ctx.drawImage(image, 0, 0, image.width, image.height, pad, pad, image.width, image.height)
return ctx.getImageData(0, 0, padded_image_dims[0], padded_image_dims[1]).data
}
Insert cell
Insert cell
Insert cell
Insert cell
kernels = ({
ones: {
dim: [5, 5],
def: Array(5 * 5).fill(1)
},
edges: {
dim: [3, 3],
def: [
0, 1, 0,
1, -4, 1,
0, 1, 0,
]
},
sobel: {
dim: [3, 3],
def: [
1, 0, -1,
2, 0, -2,
1, 0, -1,
]
},
sharr: {
dim: [3, 3],
def: [
3, 0, -3,
10, 0, -10,
3, 0, -3,
]
},
rand: {
dim: [5, 5],
def: Array(5 * 5).fill(0).map(() => {
const variance = Math.sqrt(1.0 / (5 * 5 * 3))
return util.randn(0, variance)
})
},
'vertical edge': {
dim: [3, 3],
def: [
1, 0, -1,
1, 0, -1,
1, 0, -1,
]
},
'horizontal edge': {
dim: [3, 3],
def: [
1, 1, 1,
0, 0, 0,
-1, -1, -1,
]
}
})
Insert cell
Insert cell
Insert cell
Insert cell
// original image group
imageGroup = {
const imageGroup = new THREE.Group()
scene.add(imageGroup)
return imageGroup
}
Insert cell
depthGroup = {
const depthGroup = new THREE.Group()
scene.add(depthGroup)
return depthGroup
}
Insert cell
outputGroup = {
const outputGroup = new THREE.Group()
scene.add(outputGroup)
return outputGroup
}
Insert cell
sumGroup = {
const sumGroup = new THREE.Group()
scene.add(sumGroup)
return sumGroup
}
Insert cell
Insert cell
util = ({
arrXY: (index, width) => [index % width, Math.floor(index / width)],
lerp: (min, max, x, rmin, rmax) => {
if (max - min === 0) {
return 0;
}
const norm = (x - min) / (max - min)
if (rmin === undefined || rmax === undefined) {
return norm
}
return rmin + norm * (rmax - rmin)
},
hadamard: (a, b) => a.map((x, i) => x * b[i]),
// pads a flat array of single-component color values
pad: (image, pad, image_dims, pad_value = -1) => {
let img;
if (pad) {
img = []
// left pad + image width + right pad * vertical pad
const p = (pad * 2 + image_dims[0]) * pad;

const vpad = Array(pad).fill(pad_value)
const hpad = Array(p).fill(pad_value)
// top pad
if (p) {
img.push(...hpad)
}
let j = img.length; // padded index
for (let i = 0; i < image.length; i++) {
const x = i % image_dims[0];
const y = Math.floor(i / image_dims[0]);
// left
if (x === 0) {
img.push(...vpad, image[i]);
j += pad
}
// right
else if (x === image_dims[0] - 1) {
img.push(image[i], ...vpad)
j += pad
} else {
img.push(image[i])
j++
}
}
// bottom pad
if (p) {
img.push(...hpad)
}
} else {
img = [...image]
}
return img
},
pointAt: (a, b, alpha) => {
let dir = b.clone().sub(a)
let len = dir.length()
dir = dir.normalize().multiplyScalar(len * alpha)
return a.clone().add(dir)
},
// from karpathy:
randn: (mean, variance) => {
let V1, V2, S;
do {
const U1 = Math.random();
const U2 = Math.random();
V1 = 2 * U1 - 1;
V2 = 2 * U2 - 1;
S = V1 * V1 + V2 * V2;
} while (S > 1);
let X = Math.sqrt(-2 * Math.log(S) / S) * V1;
X = mean + Math.sqrt(variance) * X;
return X;
},
})
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell
Insert cell

Purpose-built for displays of data

Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
Learn more