Spaces:
Running
Running
| import { useEffect, useMemo, useState } from "react"; | |
| export type PoolingKernel = | |
| | "Max (per channel)" | |
| | "Average (per channel)" | |
| | "Max (across channels)" | |
| | "Average (across channels)"; | |
| const MAX_SIZE = 500; | |
| function imageDataToUrl(imageData: ImageData): string { | |
| const canvas = document.createElement("canvas"); | |
| canvas.width = imageData.width; | |
| canvas.height = imageData.height; | |
| const context = canvas.getContext("2d"); | |
| if (!context) throw new Error("Failed to get canvas context"); | |
| context.putImageData(imageData, 0, 0); | |
| return canvas.toDataURL("image/png"); | |
| } | |
| function grayscale(imageData: ImageData): ImageData { | |
| const output = new ImageData(new Uint8ClampedArray(imageData.data), imageData.width, imageData.height); | |
| for (let i = 0; i < output.data.length; i += 4) { | |
| const value = 0.299 * output.data[i] + 0.587 * output.data[i + 1] + 0.114 * output.data[i + 2]; | |
| output.data[i] = output.data[i + 1] = output.data[i + 2] = value; | |
| } | |
| return output; | |
| } | |
| function pool( | |
| image: ImageData, | |
| kernelWidth: number, | |
| kernelHeight: number, | |
| stride: number, | |
| dilation: number, | |
| padding: number, | |
| kernel: PoolingKernel, | |
| ): ImageData { | |
| const effectiveWidth = dilation * (kernelWidth - 1) + 1; | |
| const effectiveHeight = dilation * (kernelHeight - 1) + 1; | |
| const width = Math.max(1, Math.floor((image.width + 2 * padding - effectiveWidth) / stride) + 1); | |
| const height = Math.max(1, Math.floor((image.height + 2 * padding - effectiveHeight) / stride) + 1); | |
| const output = new ImageData(width, height); | |
| const isMax = kernel.startsWith("Max"); | |
| const acrossChannels = kernel.endsWith("(across channels)"); | |
| for (let y = 0; y < height; y += 1) { | |
| for (let x = 0; x < width; x += 1) { | |
| const outputChannels = acrossChannels ? 1 : 3; | |
| for (let outputChannel = 0; outputChannel < outputChannels; outputChannel += 1) { | |
| let result = isMax ? -Infinity : 0; | |
| let count = 0; | |
| for (let ky = 0; ky < kernelHeight; ky += 1) { | |
| for (let kx = 0; kx < kernelWidth; kx += 1) { | |
| const sourceX = x * stride - padding + kx * dilation; | |
| const sourceY = y * stride - padding + ky * dilation; | |
| const firstChannel = acrossChannels ? 0 : outputChannel; | |
| const lastChannel = acrossChannels ? 3 : outputChannel + 1; | |
| for (let sourceChannel = firstChannel; sourceChannel < lastChannel; sourceChannel += 1) { | |
| const value = sourceX < 0 || sourceX >= image.width || sourceY < 0 || sourceY >= image.height | |
| ? 0 | |
| : image.data[(sourceY * image.width + sourceX) * 4 + sourceChannel]; | |
| result = isMax ? Math.max(result, value) : result + value; | |
| count += 1; | |
| } | |
| } | |
| } | |
| const value = isMax ? result : result / count; | |
| if (acrossChannels) { | |
| output.data[(y * width + x) * 4] = value; | |
| output.data[(y * width + x) * 4 + 1] = value; | |
| output.data[(y * width + x) * 4 + 2] = value; | |
| } else { | |
| output.data[(y * width + x) * 4 + outputChannel] = value; | |
| } | |
| } | |
| output.data[(y * width + x) * 4 + 3] = 255; | |
| } | |
| } | |
| return output; | |
| } | |
| export default function usePoolingProcessing( | |
| imageUrl: string, | |
| useColor: boolean, | |
| kernel: PoolingKernel, | |
| kernelWidth: number, | |
| kernelHeight: number, | |
| stride: number, | |
| dilation: number, | |
| padding: number, | |
| ) { | |
| const [source, setSource] = useState<ImageData | null>(null); | |
| useEffect(() => { | |
| let cancelled = false; | |
| const image = new Image(); | |
| image.crossOrigin = "anonymous"; | |
| image.onload = () => { | |
| const scale = Math.min(1, MAX_SIZE / image.naturalWidth, MAX_SIZE / image.naturalHeight); | |
| const canvas = document.createElement("canvas"); | |
| canvas.width = Math.max(1, Math.round(image.naturalWidth * scale)); | |
| canvas.height = Math.max(1, Math.round(image.naturalHeight * scale)); | |
| const context = canvas.getContext("2d"); | |
| if (!context) return; | |
| context.drawImage(image, 0, 0, canvas.width, canvas.height); | |
| if (!cancelled) setSource(context.getImageData(0, 0, canvas.width, canvas.height)); | |
| }; | |
| image.src = imageUrl; | |
| return () => { cancelled = true; }; | |
| }, [imageUrl]); | |
| return useMemo(() => { | |
| if (!source) return { input: null, output: null, inputSize: null, outputSize: null }; | |
| const input = useColor ? source : grayscale(source); | |
| const output = pool(input, kernelWidth, kernelHeight, stride, dilation, padding, kernel); | |
| return { | |
| input: imageDataToUrl(input), | |
| output: imageDataToUrl(output), | |
| inputSize: `${input.width} × ${input.height}`, | |
| outputSize: `${output.width} × ${output.height}`, | |
| }; | |
| }, [source, useColor, kernel, kernelWidth, kernelHeight, stride, dilation, padding]); | |
| } | |