Spaces:
Running
Running
File size: 4,818 Bytes
419db65 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | 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]);
}
|