Использование бинарного дерева в swift с помощью enum на примере OCR

Была цель создать приложение на mac, которое может распознавать текст кода с изображений и видео.

Хотелось сделать так, чтобы даже при большом объеме кода, текст распознавался менее, чем за секунду.

Проблема облегчается тем, что язык на котором пишут код всегда английский и ширина между всеми буквами одинаковая (моноширинный шрифт) — такие используются для программирования, и в этих шрифтах легко увидеть разницу между 1 и I, 0 и O и тд.

Если вкратце, то задача сводится к двум частям:

1. Нахождению самой буквы с ее границами

И с этим шикарно справился Vision, новый framework от Apple.

2. Распознавание буквы в заданных границах

Я решил пойти не хитрым способом и проверять определенные пиксели квадрата, в границах которого находится буквы (допустим: центр, углы, бока) и отталкиваясь от наличия или отсутствия там буквы, классифицировать что за буква.

Наглядный пример:




Как же перенести этот схематичный рисунок в код, так чтобы не закопаться в нем, и чтобы он был такой же наглядный?!

Вот здесь нам на помощь приходит бинарное дерево. Вот его каркас.

enum Tree<Node, Result> {
    ///Empty result
    case empty
    ///Result with generic type
    case r(Result)
    ///Recursive case with generic tree
    indirect case n(Node, Tree<Node, Result>, Tree<Node, Result>)
}

Теперь на основе него мы можем перенести весь наш рисунок в код.

//.c означает, если нахожу по центру пиксель с буквой, значит это "H" в противном случае "O"
//на месте .c будет находиться условие, по которому происходит ветвление
let HorOTree = TreeOCR.n(.c, .r("H"), .r("O"))

Вот как бы выглядел кусочек дерева побольше.



Можно все очень схематично разложить и легко найти нужную букву.

И последний момент, вот как выглядит сама модель, в которой происходит вся работа.

extension Tree where Node == OCROperations, Result == String {
    func find(_ colorChecker: LetterExistenceChecker, with frame: CGRect) -> String? {
        switch self {
        case .empty: return nil
        case .r(let element):
            return element
        case let .n(operation, left, right):
            let exist = operation.action(colorChecker, frame)
            return (exist ? left : right).find(colorChecker, with: frame)
        }
    }
}

В это дерево мы передаем класс LetterExistenceChecker, который отвечает за проверку наличия пикселя буквы в определенной точке в границах нужного квадрата. Конечно, я опустил множество деталей, иначе бы статья получилась слишком громоздкой. И здесь не только эти два этапа, которые были упомянуты в статье, но гораздо больше, но они были опущены, потому что цель была показать как использовать бинарное дерево и enum.

Вот демо того как работает программа, прошу заметить, что так как целью я ставил распознавание только текста с кодом, то я решил просто игнорировать весь остальной текст который не код, сделал так, чтобы прога искала только текст с кодом.


С радостью услышу ваши замечания, критику.