I’ve tried everything including ChatGPT and i just cannot get my magic wand selection tool written in swift to work. The only information I’ve found on StackOverflow has not been helpful at all. And there does not seem to be a library or code snippet out there to get me moving in the right direction.

i’m just going to post the whole class I’m working on.

struct Coordinate: Hashable {
    let x: Int
    let y: Int
}

class MagicWandSelection: NSObject {
    
    override init() {
        super.init()
    }


    func floodFill(x: Int, y: Int, width: Int, height: Int, maskData: [UInt8], checked: inout [Bool]) -> [Coordinate] {
        var stack = [Coordinate(x: x, y: y)]
        var foundPixels = [Coordinate]()

        while let point = stack.popLast() {
            let index = point.y * width + point.x
            if point.x < 0 || point.x >= width || point.y < 0 || point.y >= height || checked[index] || maskData[index] != 255 {
                continue
            }

            checked[index] = true
            foundPixels.append(point)

            let directions = [Coordinate(x: 0, y: -1), Coordinate(x: 1, y: 0), Coordinate(x: 0, y: 1), Coordinate(x: -1, y: 0)]
            for direction in directions {
                let nextPoint = Coordinate(x: point.x + direction.x, y: point.y + direction.y)
                stack.append(nextPoint)
            }
        }

        return foundPixels
    }

    
    func newMaskWithFloodFill(from point: CGPoint, in image: UIImage, selectedColor: UIColor, tolerance: Float) -> UIImage? {
        guard let cgImage = image.cgImage else { return nil }
        let width = cgImage.width
        let height = cgImage.height
        let bytesPerPixel = 4
        let bytesPerRow = width * bytesPerPixel
        let bitsPerComponent = 8

        var pixelData = [UInt8](repeating: 0, count: width * height * bytesPerPixel)
        guard let context = CGContext(data: &pixelData,
                                      width: width,
                                      height: height,
                                      bitsPerComponent: bitsPerComponent,
                                      bytesPerRow: bytesPerRow,
                                      space: CGColorSpaceCreateDeviceRGB(),
                                      bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else { return nil }

        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        let startPixelIndex = (Int(point.y) * width + Int(point.x)) * bytesPerPixel
        let startPixel = Array(pixelData[startPixelIndex..<(startPixelIndex + bytesPerPixel)])

        var maskData = [UInt8](repeating: 0, count: width * height)
        floodFill(&pixelData, &maskData, width: width, height: height, x: Int(point.x), y: Int(point.y), startPixel: startPixel, tolerance: tolerance)

        return createMaskImage(from: maskData, width: width, height: height)
    }

    func floodFill(_ pixelData: inout [UInt8], _ maskData: inout [UInt8], width: Int, height: Int, x: Int, y: Int, startPixel: [UInt8], tolerance: Float) {
        var stack = [(Int, Int)]()
        stack.append((x, y))

        while !stack.isEmpty {
            let (currentX, currentY) = stack.removeLast()
            let index = (currentY * width + currentX) * 4
            let maskIndex = currentY * width + currentX

            if maskData[maskIndex] > 0 { continue }

            let currentPixel = Array(pixelData[index..<(index + 4)])

            if !colorsMatch(pixel1: currentPixel, pixel2: startPixel, tolerance: tolerance) {
                continue
            }

            maskData[maskIndex] = 255

            if currentX > 0 {
                stack.append((currentX - 1, currentY))
            }
            if currentX < width - 1 {
                stack.append((currentX + 1, currentY))
            }
            if currentY > 0 {
                stack.append((currentX, currentY - 1))
            }
            if currentY < height - 1 {
                stack.append((currentX, currentY + 1))
            }
        }
    }

    func colorsMatch(pixel1: [UInt8], pixel2: [UInt8], tolerance: Float) -> Bool {
        let r1 = Float(pixel1[0]) / 255.0, g1 = Float(pixel1[1]) / 255.0, b1 = Float(pixel1[2]) / 255.0, a1 = Float(pixel1[3]) / 255.0
        let r2 = Float(pixel2[0]) / 255.0, g2 = Float(pixel2[1]) / 255.0, b2 = Float(pixel2[2]) / 255.0, a2 = Float(pixel2[3]) / 255.0

        let dr = r1 - r2, dg = g1 - g2, db = b1 - b2, da = a1 - a2
        let distance = sqrt(dr * dr + dg * dg + db * db + da * da)

        return distance < tolerance
    }

    func createMaskImage(from maskData: [UInt8], width: Int, height: Int) -> UIImage? {
        let bytesPerPixel = 1
        let bytesPerRow = width * bytesPerPixel
        let bitsPerComponent = 8

        var maskDataModified = maskData.map { $0 > 0 ? UInt8(255) : UInt8(0) }
        guard let context = CGContext(data: &maskDataModified,
                                      width: width,
                                      height: height,
                                      bitsPerComponent: bitsPerComponent,
                                      bytesPerRow: bytesPerRow,
                                      space: CGColorSpaceCreateDeviceGray(),
                                      bitmapInfo: CGImageAlphaInfo.none.rawValue) else { return nil }

        guard let maskCgImage = context.makeImage() else { return nil }
        return UIImage(cgImage: maskCgImage)
    }

    func createBezierPath(from maskImage: UIImage, seed: CGPoint) -> UIBezierPath? {
        guard let cgImage = maskImage.cgImage else { return nil }
        let width = Int(maskImage.size.width)
        let height = Int(maskImage.size.height)
        let bytesPerPixel = 1
        let bytesPerRow = width * bytesPerPixel
        let bitsPerComponent = 8

        var maskData = [UInt8](repeating: 0, count: width * height)
        guard let context = CGContext(data: &maskData,
                                      width: width,
                                      height: height,
                                      bitsPerComponent: bitsPerComponent,
                                      bytesPerRow: bytesPerRow,
                                      space: CGColorSpaceCreateDeviceGray(),
                                      bitmapInfo: CGImageAlphaInfo.none.rawValue) else { return nil }

        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        let path = UIBezierPath()
        let seedPoint = Coordinate(x: Int(seed.x), y: Int(seed.y))

        // Ensure the seed point is within bounds and on a white pixel
        guard seedPoint.x >= 0 && seedPoint.x < width && seedPoint.y >= 0 && seedPoint.y < height else { return nil }
        guard maskData[seedPoint.y * width + seedPoint.x] == 255 else { return nil }

        var currentPoint = seedPoint
        let directions = [
            Coordinate(x: 0, y: -1), // Up
            Coordinate(x: 1, y: 0),  // Right
            Coordinate(x: 0, y: 1),  // Down
            Coordinate(x: -1, y: 0)  // Left
        ]

        // Find the starting edge pixel by moving up and left from the seed pixel
        var startPoint: Coordinate? = nil
        for direction in [Coordinate(x: 0, y: -1), Coordinate(x: -1, y: 0)] {
            var testPoint = seedPoint
            while testPoint.x >= 0 && testPoint.y >= 0 {
                if maskData[testPoint.y * width + testPoint.x] == 255 && isEdgePixel(x: testPoint.x, y: testPoint.y, width: width, height: height, maskData: maskData) {
                    startPoint = testPoint
                    break
                }
                testPoint = Coordinate(x: testPoint.x + direction.x, y: testPoint.y + direction.y)
            }
            if startPoint != nil {
                break
            }
        }

        guard let firstPoint = startPoint else { return nil }

        currentPoint = firstPoint
        path.move(to: CGPoint(x: firstPoint.x, y: firstPoint.y))
        var visited: Set<Coordinate> = [firstPoint]

        repeat {
            var foundNext = false
            for direction in directions {
                let nextPoint = Coordinate(x: currentPoint.x + direction.x, y: currentPoint.y + direction.y)
                let nextIndex = nextPoint.y * width + nextPoint.x

                if nextPoint.x >= 0 && nextPoint.x < width && nextPoint.y >= 0 && nextPoint.y < height &&
                    maskData[nextIndex] == 255 &&
                    !visited.contains(nextPoint) &&
                    isEdgePixel(x: nextPoint.x, y: nextPoint.y, width: width, height: height, maskData: maskData) {
                    path.addLine(to: CGPoint(x: nextPoint.x, y: nextPoint.y))
                    visited.insert(nextPoint)
                    currentPoint = nextPoint
                    foundNext = true
                    break
                }
            }

            if !foundNext {
                break
            }
        } while currentPoint != firstPoint

        path.close()
        return path
    }
}


extension MagicWandSelection{
    
    func isEdgePixel(x: Int, y: Int, width: Int, height: Int, maskData: [UInt8]) -> Bool {
         if x == 0 || y == 0 || x == width - 1 || y == height - 1 {
             return true
         }
         
         let directions = [
             Coordinate(x: 0, y: -1), // Up
             Coordinate(x: 1, y: 0),  // Right
             Coordinate(x: 0, y: 1),  // Down
             Coordinate(x: -1, y: 0)  // Left
         ]
         
         for direction in directions {
             let neighborX = x + direction.x
             let neighborY = y + direction.y
             if neighborX >= 0 && neighborX < width && neighborY >= 0 && neighborY < height {
                 let neighborIndex = neighborY * width + neighborX
                 if maskData[neighborIndex] == 0 {
                     return true
                 }
             }
         }
         return false
     }

     func createMagicWandPath(from maskImage: UIImage, seed: CGPoint) -> UIBezierPath? {
         guard let cgImage = maskImage.cgImage else { return nil }
         let width = Int(maskImage.size.width)
         let height = Int(maskImage.size.height)
         let bytesPerPixel = 1
         let bytesPerRow = width * bytesPerPixel
         let bitsPerComponent = 8

         var maskData = [UInt8](repeating: 0, count: width * height)
         guard let context = CGContext(data: &maskData,
                                       width: width,
                                       height: height,
                                       bitsPerComponent: bitsPerComponent,
                                       bytesPerRow: bytesPerRow,
                                       space: CGColorSpaceCreateDeviceGray(),
                                       bitmapInfo: CGImageAlphaInfo.none.rawValue) else { return nil }

         context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

         let seedPoint = Coordinate(x: Int(seed.x), y: Int(seed.y))

         // Ensure the seed point is within bounds and on a white pixel
         guard seedPoint.x >= 0 && seedPoint.x < width && seedPoint.y >= 0 && seedPoint.y < height else { return nil }
         guard maskData[seedPoint.y * width + seedPoint.x] == 255 else { return nil }

         // Find the starting edge pixel
         guard let startPixel = findStartingEdgePixel(from: seedPoint, width: width, height: height, maskData: maskData) else {
             return nil
         }

         // Find edge pixels and construct the path
         if let edgePath = traceEdgePath(from: startPixel, width: width, height: height, maskData: maskData) {
             return edgePath
         }
         
         return nil
     }

     func findStartingEdgePixel(from seedPoint: Coordinate, width: Int, height: Int, maskData: [UInt8]) -> Coordinate? {
         for y in 0..<height {
             for x in 0..<width {
                 let index = y * width + x
                 if maskData[index] == 255 && isEdgePixel(x: x, y: y, width: width, height: height, maskData: maskData) {
                     return Coordinate(x: x, y: y)
                 }
             }
         }
         return nil
     }

     func traceEdgePath(from startPoint: Coordinate, width: Int, height: Int, maskData: [UInt8]) -> UIBezierPath? {
         let directions = [
             Coordinate(x: 0, y: -1), // Up
             Coordinate(x: 1, y: 0),  // Right
             Coordinate(x: 0, y: 1),  // Down
             Coordinate(x: -1, y: 0), // Left
             Coordinate(x: -1, y: -1), // Up-Left
             Coordinate(x: 1, y: -1),  // Up-Right
             Coordinate(x: -1, y: 1),  // Down-Left
             Coordinate(x: 1, y: 1)    // Down-Right
         ]

         let path = UIBezierPath()
         var visited = Set<Coordinate>()
         var currentPoint = startPoint
         var firstPoint = true

         repeat {
             visited.insert(currentPoint)
             if firstPoint {
                 path.move(to: CGPoint(x: currentPoint.x, y: currentPoint.y))
                 firstPoint = false
             } else {
                 path.addLine(to: CGPoint(x: currentPoint.x, y: currentPoint.y))
             }
             
             var foundNext = false
             for direction in directions {
                 let nextPoint = Coordinate(x: currentPoint.x + direction.x, y: currentPoint.y + direction.y)
                 if nextPoint.x >= 0 && nextPoint.x < width && nextPoint.y >= 0 && nextPoint.y < height &&
                    maskData[nextPoint.y * width + nextPoint.x] == 255 && !visited.contains(nextPoint) {
                     currentPoint = nextPoint
                     foundNext = true
                     break
                 }
             }
             if !foundNext {
                 break
             }
         } while currentPoint != startPoint

         path.close()
         return path
     }
    
}

I’ve tried pretty much everything I can think of including trying to use chatgpt to fix my code. What I get is a mess of a selection with the UIBezierPath zigzaging the image.
I used the ! here so that it would crash if nothing is returned this will be fixed in production code.

this is the usage example for the above code:

            let magicWand = MagicWandSelection()
            let color = UIColor.white
            let pixel = GPPixel(x: Int(pt.x), y: Int(pt.y), color: color)
            let tolerance: Float = 0.1
            let img = magicWand.newMaskWithFloodFill(from: pt, in: self.image, selectedColor: color, tolerance: tolerance)!
            self.selectionPath = magicWand.createMagicWandPath(from: img, seed: pt)

the floodFill mask returns and image of a black background with the pixels that should be outlined with a selection path as white pixels. Which should make this easy to figure out but I just cannot get past where i am.