Spaces:
Running
Running
| import { useEffect, useMemo, useState } from "react"; | |
| const MAX_SIZE: number = 500; | |
| function getBoundedDimensions(width: number, height: number): [number, number] { | |
| if (width <= MAX_SIZE && height <= MAX_SIZE) { | |
| return [width, height]; | |
| } | |
| const scale = Math.min(MAX_SIZE / width, MAX_SIZE / height); | |
| const boundedWidth = Math.max(1, Math.round(width * scale)); | |
| const boundedHeight = Math.max(1, Math.round(height * scale)); | |
| return [boundedWidth, boundedHeight]; | |
| } | |
| async function getImageData(imageUrl: string): Promise<ImageData> { | |
| const image = await new Promise<HTMLImageElement>((resolve, reject) => { | |
| const img = new Image(); | |
| img.crossOrigin = "anonymous"; | |
| img.onload = () => resolve(img); | |
| img.onerror = reject; | |
| img.src = imageUrl; | |
| }) | |
| const sourceWidth = image.naturalWidth || image.width; | |
| const sourceHeight = image.naturalHeight || image.height; | |
| const [targetWidth, targetHeight] = getBoundedDimensions(sourceWidth, sourceHeight); | |
| const canvas = document.createElement("canvas") | |
| canvas.width = targetWidth; | |
| canvas.height = targetHeight; | |
| const ctx = canvas.getContext("2d"); | |
| if (!ctx) { | |
| throw new Error("Failed to get canvas context"); | |
| } | |
| ctx.drawImage(image, 0, 0, targetWidth, targetHeight); | |
| return ctx.getImageData(0, 0, canvas.width, canvas.height); | |
| } | |
| function getImageUrl(imageData: ImageData): string { | |
| const canvas = document.createElement("canvas"); | |
| canvas.width = imageData.width; | |
| canvas.height = imageData.height; | |
| const ctx = canvas.getContext("2d") | |
| if (!ctx) { | |
| throw new Error("Failed to get canvas context"); | |
| } | |
| ctx.putImageData(imageData, 0, 0); | |
| return canvas.toDataURL("image/png"); | |
| } | |
| function convertToGrayscale(imageData: ImageData): ImageData { | |
| const output = new ImageData( | |
| new Uint8ClampedArray(imageData.data), | |
| imageData.width, | |
| imageData.height, | |
| ); | |
| const data = output.data; | |
| for (let i = 0; i < data.length; i += 4) { | |
| const r = data[i]; | |
| const g = data[i + 1]; | |
| const b = data[i + 2]; | |
| const gray = 0.299 * r + 0.587 * g + 0.114 * b; | |
| data[i] = data[i + 1] = data[i + 2] = gray; | |
| } | |
| return output; | |
| } | |
| export type ConvolutionRange = { | |
| min: number; | |
| max: number; | |
| }; | |
| type ConvolutionResult = ConvolutionRange & { | |
| values: Float32Array; | |
| width: number; | |
| height: number; | |
| channels: 1 | 3; | |
| }; | |
| function getShape(array: number[][] | number[][][]): number[] { | |
| const shape: number[] = []; | |
| let current: unknown = array; | |
| while (Array.isArray(current)) { | |
| shape.push(current.length); | |
| current = current[0]; | |
| } | |
| return shape; | |
| } | |
| function convolve( | |
| imageData: ImageData, | |
| kernel: number[][] | number[][][], | |
| useColor: boolean, | |
| stride: number, | |
| dilation: number, | |
| padding: number, | |
| ): ConvolutionResult { | |
| if (Array.isArray(kernel[0][0])) { // 3D kernel | |
| console.log('3D kernel', getShape(kernel)) | |
| return convolveColor(imageData, kernel as number[][][], stride, dilation, padding); | |
| } else { // 2D kernel | |
| console.log('2D kernel', getShape(kernel)) | |
| if (useColor) { | |
| return convolveColorDepthwise(imageData, kernel as number[][], stride, dilation, padding); | |
| } { | |
| return convolveGray(imageData, kernel as number[][], stride, dilation, padding); | |
| } | |
| } | |
| } | |
| function convolveGray(image: ImageData, kernel: number[][], stride: number, dilation: number, padding: number): ConvolutionResult { | |
| const kernelWidth = kernel[0].length; | |
| const kernelHeight = kernel.length; | |
| const width = image.width; | |
| const height = image.height; | |
| const inputData = image.data; | |
| const outputWidth = Math.max(1, Math.floor((width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1); | |
| const outputHeight = Math.max(1, Math.floor((height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1); | |
| const values = new Float32Array(outputWidth * outputHeight); | |
| let min = Infinity; | |
| let max = -Infinity; | |
| for (let y = 0; y < outputHeight; ++y) { | |
| for (let x = 0; x < outputWidth; ++x) { | |
| // dot product | |
| let sum = 0; | |
| for (let ky = 0; ky < kernelHeight; ++ky) { | |
| for (let kx = 0; kx < kernelWidth; ++kx) { | |
| const sourceX = x * stride - padding + kx * dilation; | |
| const sourceY = y * stride - padding + ky * dilation; | |
| if (sourceX < 0 || sourceX >= width || sourceY < 0 || sourceY >= height) continue; | |
| const pixelIndex = (sourceY * width + sourceX) * 4; | |
| const pixelValue = inputData[pixelIndex]; | |
| const kernelValue = kernel[ky][kx]; | |
| sum += pixelValue * kernelValue; | |
| } | |
| } | |
| values[y * outputWidth + x] = sum; | |
| min = Math.min(min, sum); | |
| max = Math.max(max, sum); | |
| } | |
| } | |
| return { values, width: outputWidth, height: outputHeight, channels: 1, min, max }; | |
| } | |
| function convolveColor(image: ImageData, kernel: number[][][], stride: number, dilation: number, padding: number): ConvolutionResult { | |
| const kernelWidth = kernel[0][0].length; | |
| const kernelHeight = kernel[0].length; | |
| const width = image.width; | |
| const height = image.height; | |
| const inputData = image.data; | |
| const outputWidth = Math.max(1, Math.floor((width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1); | |
| const outputHeight = Math.max(1, Math.floor((height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1); | |
| const values = new Float32Array(outputWidth * outputHeight); | |
| let min = Infinity; | |
| let max = -Infinity; | |
| for (let y = 0; y < outputHeight; ++y) { | |
| for (let x = 0; x < outputWidth; ++x) { | |
| // dot product over 3 channels | |
| let sum = 0; | |
| for (let ky = 0; ky < kernelHeight; ++ky) { | |
| for (let kx = 0; kx < kernelWidth; ++kx) { | |
| const sourceX = x * stride - padding + kx * dilation; | |
| const sourceY = y * stride - padding + ky * dilation; | |
| if (sourceX < 0 || sourceX >= width || sourceY < 0 || sourceY >= height) continue; | |
| const pixelIndex = (sourceY * width + sourceX) * 4; | |
| const r = inputData[pixelIndex]; | |
| const g = inputData[pixelIndex + 1]; | |
| const b = inputData[pixelIndex + 2]; | |
| const kernelR = kernel[0][ky][kx]; | |
| const kernelG = kernel[1][ky][kx]; | |
| const kernelB = kernel[2][ky][kx]; | |
| sum += r * kernelR + g * kernelG + b * kernelB; | |
| } | |
| } | |
| values[y * outputWidth + x] = sum; | |
| min = Math.min(min, sum); | |
| max = Math.max(max, sum); | |
| } | |
| } | |
| return { values, width: outputWidth, height: outputHeight, channels: 1, min, max }; | |
| } | |
| function convolveColorDepthwise(image: ImageData, kernel: number[][], stride: number, dilation: number, padding: number): ConvolutionResult { | |
| const kernelWidth = kernel[0].length; | |
| const kernelHeight = kernel.length; | |
| const outputWidth = Math.max(1, Math.floor((image.width + 2 * padding - dilation * (kernelWidth - 1) - 1) / stride) + 1); | |
| const outputHeight = Math.max(1, Math.floor((image.height + 2 * padding - dilation * (kernelHeight - 1) - 1) / stride) + 1); | |
| const values = new Float32Array(outputWidth * outputHeight * 3); | |
| let min = Infinity; | |
| let max = -Infinity; | |
| console.log('here'); | |
| for (let y = 0; y < outputHeight; ++y) { | |
| for (let x = 0; x < outputWidth; ++x) { | |
| for (let channel = 0; channel < 3; ++channel) { | |
| let sum = 0; | |
| for (let ky = 0; ky < kernelHeight; ++ky) { | |
| for (let kx = 0; kx < kernelWidth; ++kx) { | |
| const sourceX = x * stride - padding + kx * dilation; | |
| const sourceY = y * stride - padding + ky * dilation; | |
| if (sourceX < 0 || sourceX >= image.width || sourceY < 0 || sourceY >= image.height) continue; | |
| const pixelIndex = (sourceY * image.width + sourceX) * 4 + channel; | |
| sum += image.data[pixelIndex] * kernel[ky][kx]; | |
| } | |
| } | |
| values[(y * outputWidth + x) * 3 + channel] = sum; | |
| min = Math.min(min, sum); | |
| max = Math.max(max, sum); | |
| } | |
| } | |
| } | |
| console.log('done', outputHeight, outputWidth); | |
| return { values, width: outputWidth, height: outputHeight, channels: 3, min, max }; | |
| } | |
| function convolutionToImageData(result: ConvolutionResult): ImageData { | |
| const outputData = new Uint8ClampedArray(result.width * result.height * 4); | |
| const range = result.max - result.min; | |
| const pixelCount = result.width * result.height; | |
| for (let i = 0; i < pixelCount; ++i) { | |
| const outputIndex = i * 4; | |
| for (let channel = 0; channel < 3; ++channel) { | |
| const rawValue = (result.channels === 1) | |
| ? result.values[i] | |
| : result.values[i * 3 + channel]; | |
| outputData[outputIndex + channel] = (range === 0) | |
| ? 0 | |
| : 255 * (rawValue - result.min) / range; | |
| } | |
| outputData[outputIndex + 3] = 255; | |
| } | |
| console.log('conversion done'); | |
| return new ImageData(outputData, result.width, result.height); | |
| } | |
| export default function useConvolutionProcessing( | |
| rawInputImage: string, | |
| kernel: number[][][] | number[][], | |
| useColor: boolean, | |
| stride: number, | |
| dilation: number, | |
| padding: number, | |
| ): [string | null, string | null, ConvolutionRange | null] { | |
| const [rawImageData, setRawImageData] = useState<ImageData | null>(null); | |
| // extract input image data (array) | |
| useEffect(() => { | |
| let cancelled = false; | |
| async function processImage() { | |
| const imageData = await getImageData(rawInputImage); | |
| if (!cancelled) { | |
| setRawImageData(imageData); | |
| } | |
| } | |
| processImage(); | |
| return () => { | |
| cancelled = true; | |
| } | |
| }, [rawInputImage]); | |
| const processedImageData = useMemo(() => { | |
| if (!rawImageData) return null; | |
| return useColor ? rawImageData : convertToGrayscale(rawImageData); | |
| }, [rawImageData, useColor]); | |
| const convolutionResult = useMemo(() => { | |
| if (!processedImageData) return null; | |
| return convolve(processedImageData, kernel, useColor, stride, dilation, padding); | |
| }, [processedImageData, kernel, useColor, stride, dilation, padding]); | |
| const inputImage = useMemo(() => { | |
| if (!processedImageData) return null; | |
| return getImageUrl(processedImageData); | |
| }, [processedImageData]); | |
| const outputImage = useMemo(() => { | |
| if (!convolutionResult) return null; | |
| return getImageUrl(convolutionToImageData(convolutionResult)); | |
| }, [convolutionResult]); | |
| console.log('url'); | |
| const outputRange = convolutionResult | |
| ? { min: convolutionResult.min, max: convolutionResult.max } | |
| : null; | |
| console.log('range'); | |
| return [inputImage, outputImage, outputRange]; | |
| } | |