Trying to implement Hough Transform in Swift

hough.swift

//: Playground - noun: a place where people can play

import UIKit
import CoreGraphics

let image = UIImage(named: "lines.jpg")


/*
     Parameters
*/
let CLAMP_THRESHOLD = 225.0

/*
     Matrix
*/

typealias Vector = [Double]
typealias Matrix = [Vector]

fileprivate func size(ofMatrix matrix: Matrix) -> CGSize {
    if matrix.count == 0 { return CGSize.zero }
    
    let firstRow = matrix[0]
    return CGSize(width: firstRow.count, height: matrix.count)
}

fileprivate func emptyMatrix(size: CGSize) -> Matrix {
    return Matrix(
        repeating: Vector(repeating: 0, count: Int(size.width)),
        count: Int(size.height)
    )
}

fileprivate func vectorize(matrix: Matrix) -> Vector {
    let matrixSize = size(ofMatrix: matrix)
    var vector = Vector(
        repeating: 0,
        count: Int(matrixSize.width * matrixSize.height)
    )
    
    for (i, row) in matrix.enumerated() {
        for (j, _) in row.enumerated() {
            vector[i * Int(matrixSize.width) + j] = matrix[i][j]
        }
    }
    
    return vector
}

/*
     Image helpers
*/

typealias Line = (
    theta: Int,
    distance: Int,
    occurrence: Int
)

fileprivate func createImageContext(data: UnsafeMutableRawPointer, size: CGSize)
        -> CGContext {
    let width = Int(size.width)
    let height = Int(size.height)
    
    let colorspace = CGColorSpaceCreateDeviceGray()
    let context = CGContext(data: data, width: width, height: height,
                            bitsPerComponent: 8, bytesPerRow: width,
                            space: colorspace,
                            bitmapInfo: CGImageAlphaInfo.none.rawValue)
    
    return context!
}

fileprivate func makeImage(matrix: Matrix) -> UIImage {
    var vectorized = vectorize(matrix: matrix).map { UInt8($0) }
    let context = createImageContext(data: &vectorized,
                                     size: size(ofMatrix: matrix))
    
    let image = context.makeImage()
    return UIImage(cgImage: image!)
}

fileprivate func transformForRendering(houghSpace matrix: Matrix) -> Matrix {
    var maxValue = 0.0
    for (i, row) in matrix.enumerated() {
        for (j, _) in row.enumerated() {
            maxValue = max(matrix[i][j], maxValue)
        }
    }
    
    var transformedMatrix = matrix
    for (i, row) in matrix.enumerated() {
        for (j, _) in row.enumerated() {
            transformedMatrix[i][j] = matrix[i][j] * 255.0 / maxValue
        }
    }
    
    return transformedMatrix
}

func draw(lines: [Line], inImage image: UIImage, color: UIColor) -> UIImage {
    UIGraphicsBeginImageContextWithOptions(image.size, true, 0)
    let context = UIGraphicsGetCurrentContext()!
    
    image.draw(in: CGRect(origin: CGPoint(x: 0, y: 0), size: image.size))
    color.set()
    context.setLineWidth(1)
    
    for line in lines {
        let rtheta = Double(line.theta) * Double.pi / 180
        let slope = -(cos(rtheta) / sin(rtheta))
        let intercept = Double(line.distance) * (1 / sin(rtheta))
        
        let lineEquation = { CGFloat(slope * $0 + intercept) }
        
        let startCoords = CGPoint(x: CGFloat(0), y: lineEquation(0))
        let endCoords = CGPoint(x: image.size.width,
                                y: lineEquation(Double(image.size.width)))
        
        context.move(to: startCoords)
        context.addLine(to: endCoords)
        context.strokePath()
    }
    
    return UIGraphicsGetImageFromCurrentImageContext()!
}


/*
    Hough Transform
*/

fileprivate func houghSpaceDimensions(image: UIImage) -> CGSize {
    let longestDistance = sqrt(
        pow(image.size.width, 2) + pow(image.size.height, 2)
    )
    return CGSize(width: Int(ceil(longestDistance)), height: 180)
}

fileprivate func loadGrayscaleImage(image: UIImage) -> Matrix {
    let width = Int(image.size.width)
    let height = Int(image.size.height)
    
    var bitmapData = [UInt8](repeating: 0,
                             count: width * height)
    let context = createImageContext(data: &bitmapData,
                                     size: image.size)
    
    context.draw(image.cgImage!, in: CGRect(x: 0, y: 0,
                                            width: width, height: height))
    
    var imageMatrix = emptyMatrix(size: image.size)
    
    for i in 0..<bitmapData.count {
        let row = i / Int(width)
        let col = i % width
        
        imageMatrix[row][col] = Double(bitmapData[i])
    }
    
    return imageMatrix
}

func houghSpace(image: UIImage) -> Matrix {
    let matrix = loadGrayscaleImage(image: image)
    let dimensions = houghSpaceDimensions(image: image)
    var space = emptyMatrix(size: dimensions)
    
    for (i, row) in matrix.enumerated() {
        for (j, _) in row.enumerated() {
            let intensityValue = matrix[i][j]
            
            if intensityValue >= CLAMP_THRESHOLD {
                for theta in 0..<180 {
                    let rtheta = Double(theta) * Double.pi / 180
                    let distance =
                        Double(j) * cos(rtheta) - Double(i) * sin(rtheta)
                    
                    if distance >= 1 {
                        space[theta][Int(distance)] += 1
                    }
                }
            }
        }
    }
    
    return space
}

func topLines(houghSpace space: Matrix, n: Int) -> [Line] {
    let matrixSize = size(ofMatrix: space)
    
    var lines: [Line] = [Line](
        repeating: (theta: 0, distance: 0, occurrence: 0),
        count: Int(matrixSize.width * matrixSize.height)
    )
    
    for (i, row) in space.enumerated() {
        for (j, _) in row.enumerated() {
            let occurrence = space[i][j]

            lines[i * Int(matrixSize.width) + j]  = (
                theta: i,
                distance: j,
                occurrence: Int(occurrence)
            )
        }
    }
    
    let sortedLines = lines.sorted() { a, b in
        return a.occurrence < b.occurrence
    }
    
    return Array<Line>(sortedLines.suffix(n).reversed())
}

let hough = houghSpace(image: image)
let lines = topLines(houghSpace: hough, n: 5)
let imageWithLines = draw(lines: lines, inImage: image, color: .green)