Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | 428x | // TODO: add support for WGSL files
const shader = `
const MAX_STRENGTH = 65535f;
// Workgroup size - X*Y*Z must be multiple of 32 for better performance
override workGroupSizeX = 1u;
override workGroupSizeY = 1u;
override workGroupSizeZ = 1u;
// Compare the current voxel to neighbors using a 9x9x9 window
override windowSize = 9i;
struct Params {
size: vec3u,
iteration: u32,
}
// New structure to track bounds of modified voxels
struct Bounds {
minX: atomic<i32>,
minY: atomic<i32>,
minZ: atomic<i32>,
maxX: atomic<i32>,
maxY: atomic<i32>,
maxZ: atomic<i32>,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage> volumePixelData: array<f32>;
@group(0) @binding(2) var<storage, read_write> labelmap: array<u32>;
@group(0) @binding(3) var<storage, read_write> strengthData: array<f32>;
@group(0) @binding(4) var<storage> prevLabelmap: array<u32>;
@group(0) @binding(5) var<storage> prevStrengthData: array<f32>;
@group(0) @binding(6) var<storage, read_write> updatedVoxelsCounter: array<atomic<u32>>;
@group(0) @binding(7) var<storage, read_write> modifiedBounds: Bounds;
fn getPixelIndex(ijkPos: vec3u) -> u32 {
let numPixelsPerSlice = params.size.x * params.size.y;
return ijkPos.x + ijkPos.y * params.size.x + ijkPos.z * numPixelsPerSlice;
}
fn updateBounds(position: vec3i) {
// Atomically update min bounds (use min operation)
let oldMinX = atomicMin(&modifiedBounds.minX, position.x);
let oldMinY = atomicMin(&modifiedBounds.minY, position.y);
let oldMinZ = atomicMin(&modifiedBounds.minZ, position.z);
// Atomically update max bounds (use max operation)
let oldMaxX = atomicMax(&modifiedBounds.maxX, position.x);
let oldMaxY = atomicMax(&modifiedBounds.maxY, position.y);
let oldMaxZ = atomicMax(&modifiedBounds.maxZ, position.z);
}
@compute @workgroup_size(workGroupSizeX, workGroupSizeY, workGroupSizeZ)
fn main(
@builtin(global_invocation_id) globalId: vec3u,
) {
// Make sure it will not get out of bounds for volume with sizes that
// are not multiple of workGroupSize
if (
globalId.x >= params.size.x ||
globalId.y >= params.size.y ||
globalId.z >= params.size.z
) {
return;
}
// Initialize bounds for the first iteration
if (params.iteration == 0 && globalId.x == 0 && globalId.y == 0 && globalId.z == 0) {
// Initialize to opposite extremes to ensure any update will improve the bounds
atomicStore(&modifiedBounds.minX, i32(params.size.x));
atomicStore(&modifiedBounds.minY, i32(params.size.y));
atomicStore(&modifiedBounds.minZ, i32(params.size.z));
atomicStore(&modifiedBounds.maxX, -1);
atomicStore(&modifiedBounds.maxY, -1);
atomicStore(&modifiedBounds.maxZ, -1);
}
let currentCoord = vec3i(globalId);
let currentPixelIndex = getPixelIndex(globalId);
let numPixels = arrayLength(&volumePixelData);
let currentPixelValue = volumePixelData[currentPixelIndex];
if (params.iteration == 0) {
// All non-zero initial labels are given maximum strength
strengthData[currentPixelIndex] = select(MAX_STRENGTH, 0., labelmap[currentPixelIndex] == 0);
// Update bounds for non-zero initial labels
if (labelmap[currentPixelIndex] != 0) {
updateBounds(currentCoord);
}
return;
}
// It should at least copy the values from previous state
var newLabel = prevLabelmap[currentPixelIndex];
var newStrength = prevStrengthData[currentPixelIndex];
let window = i32(ceil(f32(windowSize - 1) * .5));
let minWindow = -1i * window;
let maxWindow = 1i * window;
for (var k = minWindow; k <= maxWindow; k++) {
for (var j = minWindow; j <= maxWindow; j++) {
for (var i = minWindow; i <= maxWindow; i++) {
// Skip current voxel
if (i == 0 && j == 0 && k == 0) {
continue;
}
let neighborCoord = currentCoord + vec3i(i, j, k);
// Boundary conditions. Do not grow outside of the volume
if (
neighborCoord.x < 0i || neighborCoord.x >= i32(params.size.x) ||
neighborCoord.y < 0i || neighborCoord.y >= i32(params.size.y) ||
neighborCoord.z < 0i || neighborCoord.z >= i32(params.size.z)
) {
continue;
}
let neighborIndex = getPixelIndex(vec3u(neighborCoord));
let neighborPixelValue = volumePixelData[neighborIndex];
let prevNeighborStrength = prevStrengthData[neighborIndex];
let strengthCost = abs(neighborPixelValue - currentPixelValue);
let takeoverStrength = prevNeighborStrength - strengthCost;
if (takeoverStrength > newStrength) {
newLabel = prevLabelmap[neighborIndex];
newStrength = takeoverStrength;
}
}
}
}
if (labelmap[currentPixelIndex] != newLabel) {
atomicAdd(&updatedVoxelsCounter[params.iteration], 1u);
// Update bounds for modified voxels
updateBounds(currentCoord);
}
labelmap[currentPixelIndex] = newLabel;
strengthData[currentPixelIndex] = newStrength;
}
`;
export default shader;
|