diff options
| -rw-r--r-- | Sources/PixelSort.swift | 171 | ||||
| -rw-r--r-- | Sources/Shaders.metal | 258 | ||||
| -rw-r--r-- | input.png | bin | 0 -> 45495 bytes | |||
| -rw-r--r-- | output.png | bin | 0 -> 424216 bytes |
4 files changed, 280 insertions, 149 deletions
diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift index dff52f2..fa1bee6 100644 --- a/Sources/PixelSort.swift +++ b/Sources/PixelSort.swift @@ -24,9 +24,18 @@ struct PixelSort: ParsableCommand { @Option(name: .shortAndLong, help: "Upper brightness threshold (0.0–1.0)") var upper: Float = 0.9 - @Flag(name: .shortAndLong, help: "Sort descending") + @Flag(name: .shortAndLong, help: "Sort descending (reverse)") var descending: Bool = false + @Option(name: .shortAndLong, help: "Gamma applied to sorted pixels (Unity-style composite)") + var gamma: Float = 1.0 + + @Option(name: .shortAndLong, help: "Clamp maximum sortable span length (default: image width)") + var maxSpan: Int? + + @Flag(name: .long, help: "Invert the threshold mask") + var invertMask: Bool = false + mutating func run() throws { let device = MTLCreateSystemDefaultDevice()! let compute = try Compute(device: device) @@ -44,18 +53,51 @@ struct PixelSort: ParsableCommand { let width = cgImage.width let height = cgImage.height - let desc = MTLTextureDescriptor.texture2DDescriptor( + let rgbaDesc = MTLTextureDescriptor.texture2DDescriptor( pixelFormat: .rgba8Unorm, width: width, height: height, mipmapped: false ) - desc.usage = [.shaderRead, .shaderWrite] + rgbaDesc.usage = [.shaderRead, .shaderWrite] + + let texA = device.makeTexture(descriptor: rgbaDesc)! + let texB = device.makeTexture(descriptor: rgbaDesc)! + texA.label = "original" + texB.label = "output" - let texA = device.makeTexture(descriptor: desc)! - let texB = device.makeTexture(descriptor: desc)! - texA.label = "texA" - texB.label = "texB" + let sortedTex = device.makeTexture(descriptor: rgbaDesc)! + sortedTex.label = "sorted" + + let maskDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r8Uint, + width: width, + height: height, + mipmapped: false + ) + maskDesc.usage = [.shaderRead, .shaderWrite] + let maskTex = device.makeTexture(descriptor: maskDesc)! + maskTex.label = "mask" + + let spanDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r32Uint, + width: width, + height: height, + mipmapped: false + ) + spanDesc.usage = [.shaderRead, .shaderWrite] + let spanTex = device.makeTexture(descriptor: spanDesc)! + spanTex.label = "spans" + + let valDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r16Float, + width: width, + height: height, + mipmapped: false + ) + valDesc.usage = [.shaderRead, .shaderWrite] + let valTex = device.makeTexture(descriptor: valDesc)! + valTex.label = "sortValues" // Upload pixel data let bytesPerPixel = 4 @@ -86,58 +128,68 @@ struct PixelSort: ParsableCommand { let shaderSource = try String(contentsOf: Bundle.module.url(forResource: "Shaders", withExtension: "metal")!, encoding: .utf8) let library = ShaderLibrary.source(shaderSource) - var copyPipeline = try compute.makePipeline(function: library.copyTexture) - var sortPipeline = try compute.makePipeline(function: library.bitonicSortStep) + var createMaskPipeline = try compute.makePipeline(function: library.createMask) + var clearSpanPipeline = try compute.makePipeline(function: library.clearSpanBuffer) + var identifySpansPipeline = try compute.makePipeline(function: library.identifySpans) + var rgbToValPipeline = try compute.makePipeline(function: library.rgbToSortValue) + var pixelSortPipeline = try compute.makePipeline(function: library.pixelSortSpan) + var compositePipeline = try compute.makePipeline(function: library.composite) - // Bitonic sort requires log2(nextPow2(width)) outer passes - let n = nextPowerOf2(width) - - // Ping-pong between texA and texB - var readTex = texA - var writeTex = texB - - var blockSize: Int = 2 - while blockSize <= n { - var subBlockSize = blockSize - while subBlockSize >= 2 { - // Copy readTex → writeTex so untouched pixels carry forward - copyPipeline.arguments.src = .texture(readTex) - copyPipeline.arguments.dst = .texture(writeTex) - try compute.run(pipeline: copyPipeline, width: width, height: height) + var params = Params( + width: UInt32(width), + height: UInt32(height), + sortKey: UInt32(key.metalValue), + lowerThreshold: lower, + upperThreshold: upper, + reverseSorting: descending ? 1 : 0, + gamma: gamma, + maxSpanLength: UInt32(maxSpan ?? width), + invertMask: invertMask ? 1 : 0 + ) + let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout<Params>.stride, options: .storageModeShared)! - // Run the bitonic compare-swap step - var params = BitonicParams( - width: UInt32(width), - height: UInt32(height), - blockSize: UInt32(blockSize), - subBlockSize: UInt32(subBlockSize), - sortKey: UInt32(key.metalValue), - lowerThreshold: lower, - upperThreshold: upper, - descending: descending ? 1 : 0 - ) + // 1) Mask + createMaskPipeline.arguments.colorTex = .texture(texA) + createMaskPipeline.arguments.maskTex = .texture(maskTex) + createMaskPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: createMaskPipeline, width: width, height: height) - sortPipeline.arguments.inputTexture = .texture(readTex) - sortPipeline.arguments.outputTexture = .texture(writeTex) + // 2) Clear span buffer + clearSpanPipeline.arguments.spanTex = .texture(spanTex) + clearSpanPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: clearSpanPipeline, width: width, height: height) - let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout<BitonicParams>.stride, options: .storageModeShared)! - sortPipeline.arguments.params = .buffer(paramBuffer) + // 3) Identify spans (1 thread per row: dispatch width=1) + identifySpansPipeline.arguments.maskTex = .texture(maskTex) + identifySpansPipeline.arguments.spanTex = .texture(spanTex) + identifySpansPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: identifySpansPipeline, width: 1, height: height) - try compute.run(pipeline: sortPipeline, width: width / 2, height: height) + // 4) Sort values + rgbToValPipeline.arguments.colorTex = .texture(texA) + rgbToValPipeline.arguments.valTex = .texture(valTex) + rgbToValPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: rgbToValPipeline, width: width, height: height) - // Swap - let tmp = readTex - readTex = writeTex - writeTex = tmp + // 5) Sort each span into sortedTex + pixelSortPipeline.arguments.colorTex = .texture(texA) + pixelSortPipeline.arguments.valTex = .texture(valTex) + pixelSortPipeline.arguments.spanTex = .texture(spanTex) + pixelSortPipeline.arguments.sortedTex = .texture(sortedTex) + pixelSortPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: pixelSortPipeline, width: width, height: height) - subBlockSize /= 2 - } - blockSize *= 2 - } + // 6) Composite only masked pixels into output texB + compositePipeline.arguments.maskTex = .texture(maskTex) + compositePipeline.arguments.sortedTex = .texture(sortedTex) + compositePipeline.arguments.originalTex = .texture(texA) + compositePipeline.arguments.outTex = .texture(texB) + compositePipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: compositePipeline, width: width, height: height) - // Read back from readTex (the last write destination after swap) + // Read back from output var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) - readTex.getBytes( + texB.getBytes( &outputData, bytesPerRow: bytesPerRow, from: MTLRegionMake2D(0, 0, width, height), @@ -173,25 +225,16 @@ struct PixelSort: ParsableCommand { // MARK: - Helpers -struct BitonicParams { +struct Params { var width: UInt32 var height: UInt32 - var blockSize: UInt32 - var subBlockSize: UInt32 var sortKey: UInt32 var lowerThreshold: Float var upperThreshold: Float - var descending: UInt32 -} - -func nextPowerOf2(_ n: Int) -> Int { - var v = n - 1 - v |= v >> 1 - v |= v >> 2 - v |= v >> 4 - v |= v >> 8 - v |= v >> 16 - return v + 1 + var reverseSorting: UInt32 + var gamma: Float + var maxSpanLength: UInt32 + var invertMask: UInt32 } enum SortKeyOption: String, ExpressibleByArgument, CaseIterable { diff --git a/Sources/Shaders.metal b/Sources/Shaders.metal index 493d317..31608a6 100644 --- a/Sources/Shaders.metal +++ b/Sources/Shaders.metal @@ -2,13 +2,18 @@ using namespace metal; // -------------------------------------------------------------------- -// Bitonic-sort pixel-sorting kernel for glitch art +// Span-based pixel sorting (Unity Pixel-Sorting port) // -// Each row of the image is treated as an independent array. -// Pixels whose brightness falls within [lower, upper] form a "sortable -// mask." The kernel runs successive bitonic merge passes over every -// row; masked-out pixels stay in place while masked-in pixels are -// compared-and-swapped by their sort key (brightness / hue / etc.). +// Pipeline: +// 1) createMask: mark pixels whose luminance is within thresholds +// 2) clearSpanBuffer +// 3) identifySpans: for each row, write span length at span start pixel +// 4) rgbToSortValue: compute per-pixel sort value (R/G/B/L/S/H) +// 5) pixelSortSpan: for each span start, sort pixels within the span +// 6) composite: apply sorted pixels only where mask==1 +// +// This matches the semantics of the reference Unity compute shader: +// contiguous masked regions are sorted independently. // -------------------------------------------------------------------- enum SortKey : uint { @@ -22,11 +27,11 @@ enum SortKey : uint { // ---- helpers ------------------------------------------------------- -static inline float brightness(float4 c) { - return dot(c.rgb, float3(0.2126, 0.7152, 0.0722)); +static inline float luminance(float3 rgb) { + return dot(rgb, float3(0.299, 0.587, 0.114)); } -static inline float hue(float4 c) { +static inline float hue(float3 c) { float cmax = max(c.r, max(c.g, c.b)); float cmin = min(c.r, min(c.g, c.b)); float delta = cmax - cmin; @@ -40,115 +45,198 @@ static inline float hue(float4 c) { return h; } -static inline float saturation(float4 c) { +static inline float saturation(float3 c) { float cmax = max(c.r, max(c.g, c.b)); float cmin = min(c.r, min(c.g, c.b)); if (cmax < 1e-6) return 0.0; return (cmax - cmin) / cmax; } -static inline float sort_value(float4 c, uint key) { +static inline float sort_value(float3 rgb, uint key) { switch (SortKey(key)) { - case SortKey::Brightness: return brightness(c); - case SortKey::Hue: return hue(c); - case SortKey::Saturation: return saturation(c); - case SortKey::Red: return c.r; - case SortKey::Green: return c.g; - case SortKey::Blue: return c.b; + case SortKey::Brightness: return luminance(rgb); + case SortKey::Hue: return hue(rgb); + case SortKey::Saturation: return saturation(rgb); + case SortKey::Red: return rgb.r; + case SortKey::Green: return rgb.g; + case SortKey::Blue: return rgb.b; } - return brightness(c); + return luminance(rgb); } // ---- parameters passed from the CPU side --------------------------- struct Params { - uint width; // image width (= row length) - uint height; // image height (= number of rows) - uint blockSize; // bitonic block size (power of 2, doubles each outer pass) - uint subBlockSize; // comparison distance (power of 2, halves each inner pass) - uint sortKey; // which channel to sort by (see SortKey enum) - float lowerThreshold; // brightness lower bound for mask - float upperThreshold; // brightness upper bound for mask - uint descending; // 0 = ascending, 1 = descending + uint width; + uint height; + uint sortKey; // SortKey enum + float lowerThreshold; // luminance low + float upperThreshold; // luminance high + uint reverseSorting; // 0 = normal, 1 = reverse + float gamma; // output gamma (Unity applies pow(abs(sorted), gamma)) + uint maxSpanLength; // clamp span length (safety) + uint invertMask; // 0/1 }; -// ---- bitonic compare-and-swap kernel ------------------------------- -// -// Dispatch with threads = (width/2) * height. -// Each thread handles one compare-swap pair in the current sub-pass. +// ---- create mask (0/1) --------------------------------------------- -kernel void bitonicSortStep( - texture2d<float, access::read> inputTexture [[texture(0)]], - texture2d<float, access::write> outputTexture [[texture(1)]], - constant Params ¶ms [[buffer(0)]], - uint2 gid [[thread_position_in_grid]] +kernel void createMask( + texture2d<float, access::read> colorTex [[texture(0)]], + texture2d<uint, access::write> maskTex [[texture(1)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] ) { - uint pairIndex = gid.x; // which pair within this row - uint row = gid.y; + if (gid.x >= params.width || gid.y >= params.height) return; + float3 rgb = saturate(colorTex.read(gid).rgb); + float l = luminance(rgb); + bool inRange = (l >= params.lowerThreshold) && (l <= params.upperThreshold); + uint m = inRange ? 1u : 0u; + if (params.invertMask) m = 1u - m; + maskTex.write(m, gid); +} + +// ---- clear span buffer --------------------------------------------- + +kernel void clearSpanBuffer( + texture2d<uint, access::write> spanTex [[texture(0)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] +) { + if (gid.x >= params.width || gid.y >= params.height) return; + spanTex.write(0u, gid); +} - if (row >= params.height) return; - if (pairIndex >= params.width / 2) return; +// ---- identify spans (horizontal only) ------------------------------ +// One thread per row: write span length at each span start. - // Determine the two indices to compare. - uint blockSize = params.blockSize; - uint subBlockSize = params.subBlockSize; +kernel void identifySpans( + texture2d<uint, access::read> maskTex [[texture(0)]], + texture2d<uint, access::write> spanTex [[texture(1)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] +) { + uint row = gid.y; + if (gid.x != 0 || row >= params.height) return; - // Position within block and sub-block - uint blockIndex = pairIndex / (subBlockSize / 2); - uint offset = pairIndex % (subBlockSize / 2); + uint pos = 0; + uint spanStart = 0; + uint spanLength = 0; + uint spanLimit = max(1u, params.maxSpanLength); - uint leftIdx = blockIndex * subBlockSize + offset; - uint rightIdx = leftIdx + subBlockSize / 2; + while (pos < params.width) { + uint m = maskTex.read(uint2(pos, row)).x; + pos += 1; - if (leftIdx >= params.width || rightIdx >= params.width) { - // Out of bounds — copy left pixel through unchanged. - if (leftIdx < params.width) { - outputTexture.write(inputTexture.read(uint2(leftIdx, row)), uint2(leftIdx, row)); + if (m == 0 || spanLength >= spanLimit) { + // Write at span start. Mirror Unity behavior: if we hit an unmasked pixel, + // spanLength is current count; if we hit limit while still masked, include current pixel. + if (spanLength != 0) { + uint outLen = (m == 1u) ? (spanLength + 1u) : spanLength; + spanTex.write(outLen, uint2(spanStart, row)); + } + spanStart = pos; + spanLength = 0; + } else { + spanLength += 1; } - return; } - float4 leftPixel = inputTexture.read(uint2(leftIdx, row)); - float4 rightPixel = inputTexture.read(uint2(rightIdx, row)); + if (spanLength != 0 && spanStart < params.width) { + spanTex.write(spanLength, uint2(spanStart, row)); + } +} - // Threshold mask: only sort pixels whose brightness is in range. - float leftBri = brightness(leftPixel); - float rightBri = brightness(rightPixel); - bool leftIn = leftBri >= params.lowerThreshold && leftBri <= params.upperThreshold; - bool rightIn = rightBri >= params.lowerThreshold && rightBri <= params.upperThreshold; +// ---- per-pixel sort value ------------------------------------------ - if (leftIn && rightIn) { - float leftVal = sort_value(leftPixel, params.sortKey); - float rightVal = sort_value(rightPixel, params.sortKey); +kernel void rgbToSortValue( + texture2d<float, access::read> colorTex [[texture(0)]], + texture2d<half, access::write> valTex [[texture(1)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] +) { + if (gid.x >= params.width || gid.y >= params.height) return; + float3 rgb = saturate(colorTex.read(gid).rgb); + float v = sort_value(rgb, params.sortKey); + valTex.write(half(v), gid); +} - // Direction: ascending within even blocks, descending within odd - // (standard bitonic pattern), then flip if user wants descending. - bool ascending = ((leftIdx / blockSize) % 2 == 0); - if (params.descending) ascending = !ascending; +// ---- sort pixels within a span ------------------------------------- +// One thread per pixel, but only span starts do work. +// Writes sorted pixels into `sortedTex` for the span region. - bool doSwap = ascending ? (leftVal > rightVal) : (leftVal < rightVal); - if (doSwap) { - float4 tmp = leftPixel; - leftPixel = rightPixel; - rightPixel = tmp; - } +constant uint MAX_LOCAL_SPAN = 2048; + +kernel void pixelSortSpan( + texture2d<float, access::read> colorTex [[texture(0)]], + texture2d<half, access::read> valTex [[texture(1)]], + texture2d<uint, access::read> spanTex [[texture(2)]], + texture2d<float, access::write> sortedTex [[texture(3)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] +) { + if (gid.x >= params.width || gid.y >= params.height) return; + + uint spanLength = spanTex.read(gid).x; + if (spanLength == 0) return; + + spanLength = min(spanLength, params.width - gid.x); + spanLength = min(spanLength, max(1u, params.maxSpanLength)); + spanLength = min(spanLength, MAX_LOCAL_SPAN); + + // Cache sort values for this span. + float cache[MAX_LOCAL_SPAN]; + for (uint k = 0; k < spanLength; ++k) { + cache[k] = float(valTex.read(uint2(gid.x + k, gid.y)).x); } - outputTexture.write(leftPixel, uint2(leftIdx, row)); - outputTexture.write(rightPixel, uint2(rightIdx, row)); + float minValue = cache[0]; + float maxValue = cache[0]; + uint minIndex = 0; + uint maxIndex = 0; + + uint steps = (spanLength / 2) + 1; + for (uint i = 0; i < steps; ++i) { + for (uint j = 1; j < spanLength; ++j) { + float v = cache[j]; + // Unity checks `v == saturate(v)` to ignore sentinels; equivalent is 0..1. + if (v >= 0.0f && v <= 1.0f) { + if (v < minValue) { minValue = v; minIndex = j; } + if (maxValue < v) { maxValue = v; maxIndex = j; } + } + } + + uint dstMin = params.reverseSorting ? i : (spanLength - i - 1); + uint dstMax = params.reverseSorting ? (spanLength - i - 1) : i; + + float4 cMin = colorTex.read(uint2(gid.x + minIndex, gid.y)); + float4 cMax = colorTex.read(uint2(gid.x + maxIndex, gid.y)); + + sortedTex.write(cMin, uint2(gid.x + dstMin, gid.y)); + sortedTex.write(cMax, uint2(gid.x + dstMax, gid.y)); + + cache[minIndex] = 2.0f; + cache[maxIndex] = -2.0f; + minValue = 1.0f; + maxValue = -1.0f; + } } -// ---- simple copy kernel for pixels not touched by a compare-swap --- -// -// After each step we need untouched pixels carried forward. Instead of -// a separate pass we do a full-image copy first, then the sort step -// overwrites the pairs it touches. This kernel does that copy. +// ---- composite sorted pixels onto original ------------------------- -kernel void copyTexture( - texture2d<float, access::read> src [[texture(0)]], - texture2d<float, access::write> dst [[texture(1)]], - uint2 gid [[thread_position_in_grid]] +kernel void composite( + texture2d<uint, access::read> maskTex [[texture(0)]], + texture2d<float, access::read> sortedTex [[texture(1)]], + texture2d<float, access::read> originalTex [[texture(2)]], + texture2d<float, access::write> outTex [[texture(3)]], + constant Params ¶ms [[buffer(0)]], + uint2 gid [[thread_position_in_grid]] ) { - if (gid.x >= src.get_width() || gid.y >= src.get_height()) return; - dst.write(src.read(gid), gid); + if (gid.x >= params.width || gid.y >= params.height) return; + + float4 c = originalTex.read(gid); + if (maskTex.read(gid).x == 1u) { + float4 s = sortedTex.read(gid); + c = pow(abs(s), float4(params.gamma)); + } + outTex.write(c, gid); } diff --git a/input.png b/input.png Binary files differnew file mode 100644 index 0000000..9c9342c --- /dev/null +++ b/input.png diff --git a/output.png b/output.png Binary files differnew file mode 100644 index 0000000..a29fba9 --- /dev/null +++ b/output.png |
