Spaces:

TechAvenger
/

kernel-convolution-visualizer

Paused

kernel-convolution-visualizer / src /lib /convolution /engine.ts

mihirmithani

Phase 3

f3aa2b6 3 months ago

9.18 kB

import { PaddingMode, ConvolutionStep, ImageDataExtended } from '@/types'

/**
 * Computes a single convolution step for a specific pixel position.
 * Useful for static visualization/inspection without running the full animation.
 */
export function computeConvolutionStep(
  image: ImageDataExtended,
  kernel: number[][],
  padding: PaddingMode,
  normalize: boolean,
  x: number,
  y: number,
  channel: number
): ConvolutionStep {
  const { width, height, channels } = image
  const kernelSize = kernel.length
  const halfKernel = Math.floor(kernelSize / 2)

  const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
  const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

  const inputPatch: number[][] = []
  const multiplications: number[][] = []
  let sum = 0

  for (let ky = 0; ky < kernelSize; ky++) {
    const patchRow: number[] = []
    const multRow: number[] = []
    for (let kx = 0; kx < kernelSize; kx++) {
      const px = x + kx - halfKernel
      const py = y + ky - halfKernel
      const pixelValue = getPixel(image, px, py, channel, padding)
      const kernelValue = kernel[ky][kx]
      patchRow.push(pixelValue)
      const mult = pixelValue * kernelValue
      multRow.push(mult)
      sum += mult
    }
    inputPatch.push(patchRow)
    multiplications.push(multRow)
  }

  const outputValue = sum
  const normalizedValue = sum * normalizeFactor

  // Calculate progress (for consistency with generator)
  const totalPixels = width * height * channels
  const processedPixels = (y * width + x) * channels + channel + 1
  const progress = processedPixels / totalPixels

  return {
    x,
    y,
    channel,
    inputPatch,
    kernel,
    multiplications,
    sum: outputValue,
    outputValue: normalizedValue,
    normalizedValue,
    progress,
  }
}

export function getPixel(
  image: ImageDataExtended,
  x: number,
  y: number,
  channel: number,
  padding: PaddingMode
): number {
  const { width, height, channels, data } = image

  let px = x
  let py = y

  switch (padding) {
    case 'zero':
      if (px < 0 || px >= width || py < 0 || py >= height) return 0
      break
    case 'replicate':
      px = Math.max(0, Math.min(px, width - 1))
      py = Math.max(0, Math.min(py, height - 1))
      break
    case 'reflect':
      if (px < 0) px = -px
      if (px >= width) px = 2 * width - px - 2
      if (py < 0) py = -py
      if (py >= height) py = 2 * height - py - 2
      px = Math.max(0, Math.min(px, width - 1))
      py = Math.max(0, Math.min(py, height - 1))
      break
    case 'wrap':
      px = ((px % width) + width) % width
      py = ((py % height) + height) % height
      break
  }

  const index = (py * width + px) * channels + channel
  return data[index] / 255
}

export function setPixel(
  image: ImageDataExtended,
  x: number,
  y: number,
  channel: number,
  value: number
): void {
  const { width, channels, data } = image
  if (x < 0 || x >= width || y < 0 || y >= image.height) return
  const index = (y * width + x) * channels + channel
  data[index] = Math.max(0, Math.min(255, Math.round(value * 255)))
}

export function* convolutionSteps(
  image: ImageDataExtended,
  kernel: number[][],
  padding: PaddingMode,
  normalize: boolean,
  onProgress?: (progress: number) => void
): Generator<ConvolutionStep> {
  const { width, height, channels } = image
  const kernelSize = kernel.length
  const halfKernel = Math.floor(kernelSize / 2)

  const totalPixels = width * height * channels
  let processedPixels = 0

  const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
  const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

  const outputData = new Float32Array(width * height * channels)

  // For progressive rendering: yield partial output every N rows
  const rowsPerYield = Math.max(1, Math.floor(height / 20)) // ~20 yields per image
  let rowsSinceYield = 0

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      for (let c = 0; c < channels; c++) {
        const inputPatch: number[][] = []
        const multiplications: number[][] = []
        let sum = 0

        for (let ky = 0; ky < kernelSize; ky++) {
          const patchRow: number[] = []
          const multRow: number[] = []
          for (let kx = 0; kx < kernelSize; kx++) {
            const px = x + kx - halfKernel
            const py = y + ky - halfKernel
            const pixelValue = getPixel(image, px, py, c, padding)
            const kernelValue = kernel[ky][kx]
            patchRow.push(pixelValue)
            const mult = pixelValue * kernelValue
            multRow.push(mult)
            sum += mult
          }
          inputPatch.push(patchRow)
          multiplications.push(multRow)
        }

        const outputValue = sum
        const normalizedValue = sum * normalizeFactor

        const outputIndex = (y * width + x) * channels + c
        outputData[outputIndex] = normalizedValue

        processedPixels++
        const progress = processedPixels / totalPixels

        // Create partial output image for progressive rendering
        let partialOutputImageData: ImageData | undefined
        let isPartialOutput = false

        if (c === channels - 1 && x === width - 1) {
          // End of row - check if we should yield partial output
          rowsSinceYield++
          if (rowsSinceYield >= rowsPerYield || y === height - 1) {
            // Create partial output image
            const partialImageData = new ImageData(width, height)
            for (let i = 0; i < outputData.length; i++) {
              // Only fill up to current progress
              const pixelIdx = Math.floor(i / channels)
              const currentPixelIdx = y * width + x
              if (pixelIdx <= currentPixelIdx) {
                partialImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
              } else {
                // Unprocessed pixels - show as transparent or original
                partialImageData.data[i] = 0
                partialImageData.data[i + 3] = 0 // transparent alpha for unprocessed
              }
            }
            // Fill alpha for processed pixels
            for (let py = 0; py <= y; py++) {
              for (let px = 0; px < width; px++) {
                for (let pc = 0; pc < channels; pc++) {
                  const idx = (py * width + px) * channels + pc
                  const imgIdx = (py * width + px) * 4 + pc
                  if (pc < 3) {
                    partialImageData.data[imgIdx] = Math.max(0, Math.min(255, Math.round(outputData[idx] * 255)))
                  } else {
                    partialImageData.data[imgIdx] = 255
                  }
                }
              }
            }

            partialOutputImageData = partialImageData
            isPartialOutput = true
            rowsSinceYield = 0
          }
        }

        yield {
          x,
          y,
          channel: c,
          inputPatch,
          kernel,
          multiplications,
          sum: outputValue,
          outputValue: normalizedValue,
          normalizedValue,
          progress,
          partialOutputImageData,
          isPartialOutput,
        }

        if (onProgress) onProgress(progress)
      }
    }
  }

  // Convert output to ImageData
  const outputImageData = new ImageData(width, height)
  for (let i = 0; i < outputData.length; i++) {
    outputImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
  }

  // Final yield with complete output
  yield {
    x: -1,
    y: -1,
    channel: -1,
    inputPatch: [],
    kernel,
    multiplications: [],
    sum: 0,
    outputValue: 0,
    normalizedValue: 0,
    progress: 1,
    partialOutputImageData: outputImageData,
    isPartialOutput: true,
  }
}

export function convolveImage(
  imageData: ImageData,
  kernel: number[][],
  padding: PaddingMode = 'zero',
  normalize: boolean = true
): ImageData {
  const extended: ImageDataExtended = {
    data: imageData.data,
    width: imageData.width,
    height: imageData.height,
    channels: 4,
  }

  const { width, height } = imageData
  const kernelSize = kernel.length
  const halfKernel = Math.floor(kernelSize / 2)
  const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
  const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

  const outputData = new Uint8ClampedArray(imageData.data.length)

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      for (let c = 0; c < 3; c++) {
        let sum = 0
        for (let ky = 0; ky < kernelSize; ky++) {
          for (let kx = 0; kx < kernelSize; kx++) {
            const px = x + kx - halfKernel
            const py = y + ky - halfKernel
            const pixelValue = getPixel(extended, px, py, c, padding)
            sum += pixelValue * kernel[ky][kx]
          }
        }
        const outputIndex = (y * width + x) * 4 + c
        outputData[outputIndex] = Math.max(0, Math.min(255, Math.round(sum * normalizeFactor * 255)))
      }
      outputData[(y * width + x) * 4 + 3] = 255
    }
  }

  return new ImageData(outputData, imageData.width, imageData.height)
}

	import { PaddingMode, ConvolutionStep, ImageDataExtended } from '@/types'

	/**
	* Computes a single convolution step for a specific pixel position.
	* Useful for static visualization/inspection without running the full animation.
	*/
	export function computeConvolutionStep(
	image: ImageDataExtended,
	kernel: number[][],
	padding: PaddingMode,
	normalize: boolean,
	x: number,
	y: number,
	channel: number
	): ConvolutionStep {
	const { width, height, channels } = image
	const kernelSize = kernel.length
	const halfKernel = Math.floor(kernelSize / 2)

	const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
	const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

	const inputPatch: number[][] = []
	const multiplications: number[][] = []
	let sum = 0

	for (let ky = 0; ky < kernelSize; ky++) {
	const patchRow: number[] = []
	const multRow: number[] = []
	for (let kx = 0; kx < kernelSize; kx++) {
	const px = x + kx - halfKernel
	const py = y + ky - halfKernel
	const pixelValue = getPixel(image, px, py, channel, padding)
	const kernelValue = kernel[ky][kx]
	patchRow.push(pixelValue)
	const mult = pixelValue * kernelValue
	multRow.push(mult)
	sum += mult
	}
	inputPatch.push(patchRow)
	multiplications.push(multRow)
	}

	const outputValue = sum
	const normalizedValue = sum * normalizeFactor

	// Calculate progress (for consistency with generator)
	const totalPixels = width * height * channels
	const processedPixels = (y * width + x) * channels + channel + 1
	const progress = processedPixels / totalPixels

	return {
	x,
	y,
	channel,
	inputPatch,
	kernel,
	multiplications,
	sum: outputValue,
	outputValue: normalizedValue,
	normalizedValue,
	progress,
	}
	}

	export function getPixel(
	image: ImageDataExtended,
	x: number,
	y: number,
	channel: number,
	padding: PaddingMode
	): number {
	const { width, height, channels, data } = image

	let px = x
	let py = y

	switch (padding) {
	case 'zero':
	if (px < 0 \|\| px >= width \|\| py < 0 \|\| py >= height) return 0
	break
	case 'replicate':
	px = Math.max(0, Math.min(px, width - 1))
	py = Math.max(0, Math.min(py, height - 1))
	break
	case 'reflect':
	if (px < 0) px = -px
	if (px >= width) px = 2 * width - px - 2
	if (py < 0) py = -py
	if (py >= height) py = 2 * height - py - 2
	px = Math.max(0, Math.min(px, width - 1))
	py = Math.max(0, Math.min(py, height - 1))
	break
	case 'wrap':
	px = ((px % width) + width) % width
	py = ((py % height) + height) % height
	break
	}

	const index = (py * width + px) * channels + channel
	return data[index] / 255
	}

	export function setPixel(
	image: ImageDataExtended,
	x: number,
	y: number,
	channel: number,
	value: number
	): void {
	const { width, channels, data } = image
	if (x < 0 \|\| x >= width \|\| y < 0 \|\| y >= image.height) return
	const index = (y * width + x) * channels + channel
	data[index] = Math.max(0, Math.min(255, Math.round(value * 255)))
	}

	export function* convolutionSteps(
	image: ImageDataExtended,
	kernel: number[][],
	padding: PaddingMode,
	normalize: boolean,
	onProgress?: (progress: number) => void
	): Generator<ConvolutionStep> {
	const { width, height, channels } = image
	const kernelSize = kernel.length
	const halfKernel = Math.floor(kernelSize / 2)

	const totalPixels = width * height * channels
	let processedPixels = 0

	const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
	const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

	const outputData = new Float32Array(width * height * channels)

	// For progressive rendering: yield partial output every N rows
	const rowsPerYield = Math.max(1, Math.floor(height / 20)) // ~20 yields per image
	let rowsSinceYield = 0

	for (let y = 0; y < height; y++) {
	for (let x = 0; x < width; x++) {
	for (let c = 0; c < channels; c++) {
	const inputPatch: number[][] = []
	const multiplications: number[][] = []
	let sum = 0

	for (let ky = 0; ky < kernelSize; ky++) {
	const patchRow: number[] = []
	const multRow: number[] = []
	for (let kx = 0; kx < kernelSize; kx++) {
	const px = x + kx - halfKernel
	const py = y + ky - halfKernel
	const pixelValue = getPixel(image, px, py, c, padding)
	const kernelValue = kernel[ky][kx]
	patchRow.push(pixelValue)
	const mult = pixelValue * kernelValue
	multRow.push(mult)
	sum += mult
	}
	inputPatch.push(patchRow)
	multiplications.push(multRow)
	}

	const outputValue = sum
	const normalizedValue = sum * normalizeFactor

	const outputIndex = (y * width + x) * channels + c
	outputData[outputIndex] = normalizedValue

	processedPixels++
	const progress = processedPixels / totalPixels

	// Create partial output image for progressive rendering
	let partialOutputImageData: ImageData \| undefined
	let isPartialOutput = false

	if (c === channels - 1 && x === width - 1) {
	// End of row - check if we should yield partial output
	rowsSinceYield++
	if (rowsSinceYield >= rowsPerYield \|\| y === height - 1) {
	// Create partial output image
	const partialImageData = new ImageData(width, height)
	for (let i = 0; i < outputData.length; i++) {
	// Only fill up to current progress
	const pixelIdx = Math.floor(i / channels)
	const currentPixelIdx = y * width + x
	if (pixelIdx <= currentPixelIdx) {
	partialImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
	} else {
	// Unprocessed pixels - show as transparent or original
	partialImageData.data[i] = 0
	partialImageData.data[i + 3] = 0 // transparent alpha for unprocessed
	}
	}
	// Fill alpha for processed pixels
	for (let py = 0; py <= y; py++) {
	for (let px = 0; px < width; px++) {
	for (let pc = 0; pc < channels; pc++) {
	const idx = (py * width + px) * channels + pc
	const imgIdx = (py * width + px) * 4 + pc
	if (pc < 3) {
	partialImageData.data[imgIdx] = Math.max(0, Math.min(255, Math.round(outputData[idx] * 255)))
	} else {
	partialImageData.data[imgIdx] = 255
	}
	}
	}
	}

	partialOutputImageData = partialImageData
	isPartialOutput = true
	rowsSinceYield = 0
	}
	}

	yield {
	x,
	y,
	channel: c,
	inputPatch,
	kernel,
	multiplications,
	sum: outputValue,
	outputValue: normalizedValue,
	normalizedValue,
	progress,
	partialOutputImageData,
	isPartialOutput,
	}

	if (onProgress) onProgress(progress)
	}
	}
	}

	// Convert output to ImageData
	const outputImageData = new ImageData(width, height)
	for (let i = 0; i < outputData.length; i++) {
	outputImageData.data[i] = Math.max(0, Math.min(255, Math.round(outputData[i] * 255)))
	}

	// Final yield with complete output
	yield {
	x: -1,
	y: -1,
	channel: -1,
	inputPatch: [],
	kernel,
	multiplications: [],
	sum: 0,
	outputValue: 0,
	normalizedValue: 0,
	progress: 1,
	partialOutputImageData: outputImageData,
	isPartialOutput: true,
	}
	}

	export function convolveImage(
	imageData: ImageData,
	kernel: number[][],
	padding: PaddingMode = 'zero',
	normalize: boolean = true
	): ImageData {
	const extended: ImageDataExtended = {
	data: imageData.data,
	width: imageData.width,
	height: imageData.height,
	channels: 4,
	}

	const { width, height } = imageData
	const kernelSize = kernel.length
	const halfKernel = Math.floor(kernelSize / 2)
	const kernelSum = kernel.flat().reduce((a, b) => a + b, 0)
	const normalizeFactor = normalize && kernelSum !== 0 ? 1 / kernelSum : 1

	const outputData = new Uint8ClampedArray(imageData.data.length)

	for (let y = 0; y < height; y++) {
	for (let x = 0; x < width; x++) {
	for (let c = 0; c < 3; c++) {
	let sum = 0
	for (let ky = 0; ky < kernelSize; ky++) {
	for (let kx = 0; kx < kernelSize; kx++) {
	const px = x + kx - halfKernel
	const py = y + ky - halfKernel
	const pixelValue = getPixel(extended, px, py, c, padding)
	sum += pixelValue * kernel[ky][kx]
	}
	}
	const outputIndex = (y * width + x) * 4 + c
	outputData[outputIndex] = Math.max(0, Math.min(255, Math.round(sum * normalizeFactor * 255)))
	}
	outputData[(y * width + x) * 4 + 3] = 255
	}
	}

	return new ImageData(outputData, imageData.width, imageData.height)
	}