From 26dee10dcc64405b2f729efebf6e4e783b280d15 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Mon, 27 Apr 2026 14:39:13 +0300 Subject: Add initial files --- Sources/PixelSort.swift | 210 ++++++++++++++++++++++++++++++++++++++++++++++++ Sources/Shaders.metal | 154 +++++++++++++++++++++++++++++++++++ 2 files changed, 364 insertions(+) create mode 100644 Sources/PixelSort.swift create mode 100644 Sources/Shaders.metal (limited to 'Sources') diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift new file mode 100644 index 0000000..dff52f2 --- /dev/null +++ b/Sources/PixelSort.swift @@ -0,0 +1,210 @@ +import AppKit +import ArgumentParser +import Compute +import Metal + +@main +struct PixelSort: ParsableCommand { + static let configuration = CommandConfiguration( + abstract: "GPU-accelerated pixel sorting for glitch art" + ) + + @Argument(help: "Input image path") + var input: String + + @Argument(help: "Output image path") + var output: String + + @Option(name: .shortAndLong, help: "Sort key: brightness, hue, saturation, red, green, blue") + var key: SortKeyOption = .brightness + + @Option(name: .shortAndLong, help: "Lower brightness threshold (0.0–1.0)") + var lower: Float = 0.1 + + @Option(name: .shortAndLong, help: "Upper brightness threshold (0.0–1.0)") + var upper: Float = 0.9 + + @Flag(name: .shortAndLong, help: "Sort descending") + var descending: Bool = false + + mutating func run() throws { + let device = MTLCreateSystemDefaultDevice()! + let compute = try Compute(device: device) + + // Load image into a Metal texture + let inputURL = URL(fileURLWithPath: input) + let outputURL = URL(fileURLWithPath: self.output) + + guard let nsImage = NSImage(contentsOf: inputURL), + let cgImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil) + else { + throw ValidationError("Could not load image at \(input)") + } + + let width = cgImage.width + let height = cgImage.height + + let desc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .rgba8Unorm, + width: width, + height: height, + mipmapped: false + ) + desc.usage = [.shaderRead, .shaderWrite] + + let texA = device.makeTexture(descriptor: desc)! + let texB = device.makeTexture(descriptor: desc)! + texA.label = "texA" + texB.label = "texB" + + // Upload pixel data + let bytesPerPixel = 4 + let bytesPerRow = bytesPerPixel * width + var pixelData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) + let colorSpace = CGColorSpaceCreateDeviceRGB() + guard let ctx = CGContext( + data: &pixelData, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue + ) else { + throw ValidationError("Failed to create CGContext") + } + ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height)) + + texA.replace( + region: MTLRegionMake2D(0, 0, width, height), + mipmapLevel: 0, + withBytes: pixelData, + bytesPerRow: bytesPerRow + ) + + // Load shaders + 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) + + // 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) + + // 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 + ) + + sortPipeline.arguments.inputTexture = .texture(readTex) + sortPipeline.arguments.outputTexture = .texture(writeTex) + + let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout.stride, options: .storageModeShared)! + sortPipeline.arguments.params = .buffer(paramBuffer) + + try compute.run(pipeline: sortPipeline, width: width / 2, height: height) + + // Swap + let tmp = readTex + readTex = writeTex + writeTex = tmp + + subBlockSize /= 2 + } + blockSize *= 2 + } + + // Read back from readTex (the last write destination after swap) + var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) + readTex.getBytes( + &outputData, + bytesPerRow: bytesPerRow, + from: MTLRegionMake2D(0, 0, width, height), + mipmapLevel: 0 + ) + + // Save output + guard let outCtx = CGContext( + data: &outputData, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue + ), + let outCGImage = outCtx.makeImage() + else { + throw ValidationError("Failed to create output image") + } + + let nsOutImage = NSBitmapImageRep(cgImage: outCGImage) + let ext = outputURL.pathExtension.lowercased() + let fileType: NSBitmapImageRep.FileType = ext == "jpg" || ext == "jpeg" ? .jpeg : .png + guard let data = nsOutImage.representation(using: fileType, properties: [:]) else { + throw ValidationError("Failed to encode output image") + } + try data.write(to: outputURL) + + print("Pixel-sorted image saved to \(self.output) (\(width)×\(height))") + } +} + +// MARK: - Helpers + +struct BitonicParams { + 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 +} + +enum SortKeyOption: String, ExpressibleByArgument, CaseIterable { + case brightness, hue, saturation, red, green, blue + + var metalValue: Int { + switch self { + case .brightness: return 0 + case .hue: return 1 + case .saturation: return 2 + case .red: return 3 + case .green: return 4 + case .blue: return 5 + } + } +} diff --git a/Sources/Shaders.metal b/Sources/Shaders.metal new file mode 100644 index 0000000..493d317 --- /dev/null +++ b/Sources/Shaders.metal @@ -0,0 +1,154 @@ +#include +using namespace metal; + +// -------------------------------------------------------------------- +// Bitonic-sort pixel-sorting kernel for glitch art +// +// 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.). +// -------------------------------------------------------------------- + +enum SortKey : uint { + Brightness = 0, + Hue = 1, + Saturation = 2, + Red = 3, + Green = 4, + Blue = 5, +}; + +// ---- helpers ------------------------------------------------------- + +static inline float brightness(float4 c) { + return dot(c.rgb, float3(0.2126, 0.7152, 0.0722)); +} + +static inline float hue(float4 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; + if (delta < 1e-6) return 0.0; + float h; + if (cmax == c.r) h = fmod((c.g - c.b) / delta, 6.0); + else if (cmax == c.g) h = (c.b - c.r) / delta + 2.0; + else h = (c.r - c.g) / delta + 4.0; + h /= 6.0; + if (h < 0.0) h += 1.0; + return h; +} + +static inline float saturation(float4 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) { + 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; + } + return brightness(c); +} + +// ---- 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 +}; + +// ---- bitonic compare-and-swap kernel ------------------------------- +// +// Dispatch with threads = (width/2) * height. +// Each thread handles one compare-swap pair in the current sub-pass. + +kernel void bitonicSortStep( + texture2d inputTexture [[texture(0)]], + texture2d outputTexture [[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 (row >= params.height) return; + if (pairIndex >= params.width / 2) return; + + // Determine the two indices to compare. + uint blockSize = params.blockSize; + uint subBlockSize = params.subBlockSize; + + // Position within block and sub-block + uint blockIndex = pairIndex / (subBlockSize / 2); + uint offset = pairIndex % (subBlockSize / 2); + + uint leftIdx = blockIndex * subBlockSize + offset; + uint rightIdx = leftIdx + subBlockSize / 2; + + 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)); + } + return; + } + + float4 leftPixel = inputTexture.read(uint2(leftIdx, row)); + float4 rightPixel = inputTexture.read(uint2(rightIdx, 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; + + if (leftIn && rightIn) { + float leftVal = sort_value(leftPixel, params.sortKey); + float rightVal = sort_value(rightPixel, params.sortKey); + + // 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; + + bool doSwap = ascending ? (leftVal > rightVal) : (leftVal < rightVal); + if (doSwap) { + float4 tmp = leftPixel; + leftPixel = rightPixel; + rightPixel = tmp; + } + } + + outputTexture.write(leftPixel, uint2(leftIdx, row)); + outputTexture.write(rightPixel, uint2(rightIdx, row)); +} + +// ---- 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. + +kernel void copyTexture( + texture2d src [[texture(0)]], + texture2d dst [[texture(1)]], + uint2 gid [[thread_position_in_grid]] +) { + if (gid.x >= src.get_width() || gid.y >= src.get_height()) return; + dst.write(src.read(gid), gid); +} -- cgit v1.2.3