PDF_Office/PDF Office/PDF Master/Class/PDFTools/OCRNew/Model/KMGOCRManagerNew.swift

//
//  KMGOCRManagerNew.swift
//  PDF Master
//
//  Created by liujiajie on 2023/11/15.
//

import Foundation
import Cocoa
import Vision


let KMGOCRLanguageCodeKey   = "KMGOCRLanguageCodeKey"
let KMGOCRLanguageStringKey = "KMGOCRLanguageStringKey"
let KMImageScale = 4.0

@objc enum KMOCRType: Int {
    case Google = 0
    case Apple
}

@objc(KMGOCRManagerNewDelegate)
protocol KMGOCRManagerNewDelegate: AnyObject {
    @objc optional func GOCRManagerDidStartOCR(_ manager: KMGOCRManagerNew)
    @objc optional func GOCRManagerDidFinishOCR(_ manager: KMGOCRManagerNew)
    @objc optional func GOCRManagerDidCancel(_ manager:KMGOCRManagerNew, atIndex:Int)
    @objc optional func GOCRManagerDidStart(_ manager:KMGOCRManagerNew, atIndex:Int)
    @objc optional func GOCRManagerDidFinish(_ manager:KMGOCRManagerNew, atIndex:Int, results: [Any])
    @objc optional func GOCRManagerDidFail(_ manager:KMGOCRManagerNew, atIndex:Int, error: Error?)
}

class KMOCROperationQueue: OperationQueue{
    static let sharedInstance: KMOCROperationQueue = {
        let queue = KMOCROperationQueue()
        return queue
    }()
    
    func addOCROperation(op: Operation) {
        self.addOperation(op)
    }
    func cancelAll() {
        self.cancelAllOperations()
    }
    
}

@objcMembers class KMGOCRManagerNew: NSObject, KMGOCROperationDelegate{
    var delegate: KMGOCRManagerNewDelegate?
    var images: Array<Any>?
    var OCRType: KMOCRType = .Apple
    var selectedLanguages: Array<Any>?
    var isOCR = false
    
    var languages: Array<String>?
    var fileType: String = ""
    var ocrPath: Array<Any>?
    var filePath: URL?
    var finishIndex: Int = 0
    var appleRequest: VNRecognizeTextRequest?
    var appleRecognitionMode: VNRequestTextRecognitionLevel?
    
    override init() {
        super.init()
    }
    
    static let defaultManager: KMGOCRManagerNew = {
       let manager = KMGOCRManagerNew()
        return manager
    }()
    /*class func languages() -> [[String: Any]] {
        if KMGOCRManagerNew.defaultManager.OCRType == .Google {
            return [[KMGOCRLanguageCodeKey: "af", KMGOCRLanguageStringKey: "Afrikaans"],
                    [KMGOCRLanguageCodeKey: "sq", KMGOCRLanguageStringKey: "Albanian"],
                    [KMGOCRLanguageCodeKey: "ar", KMGOCRLanguageStringKey: "Arabic"],
                    [KMGOCRLanguageCodeKey: "hy", KMGOCRLanguageStringKey: "Armenian"],
                    [KMGOCRLanguageCodeKey: "az", KMGOCRLanguageStringKey: "Azerbaijani"],
                    [KMGOCRLanguageCodeKey: "eu", KMGOCRLanguageStringKey: "Basque"],
                    [KMGOCRLanguageCodeKey: "be", KMGOCRLanguageStringKey: "Belarusian"],
                    [KMGOCRLanguageCodeKey: "bn", KMGOCRLanguageStringKey: "Bengali"],
                    [KMGOCRLanguageCodeKey: "bs", KMGOCRLanguageStringKey: "Bosnian"],
                    [KMGOCRLanguageCodeKey: "bg", KMGOCRLanguageStringKey: "Bulgarian"],
                    [KMGOCRLanguageCodeKey: "ca", KMGOCRLanguageStringKey: "Catalan"],
                    [KMGOCRLanguageCodeKey: "ceb", KMGOCRLanguageStringKey: "Cebuano"],
                    [KMGOCRLanguageCodeKey: "ny", KMGOCRLanguageStringKey: "Chichewa"],
                    [KMGOCRLanguageCodeKey: "zh-CN", KMGOCRLanguageStringKey: "Chinese Simplified"],
                    [KMGOCRLanguageCodeKey: "zh-TW", KMGOCRLanguageStringKey: "Chinese Traditional"],
                    [KMGOCRLanguageCodeKey: "hr", KMGOCRLanguageStringKey: "Croatian"],
                    [KMGOCRLanguageCodeKey: "cs", KMGOCRLanguageStringKey: "Czech"],
                    [KMGOCRLanguageCodeKey: "da", KMGOCRLanguageStringKey: "Danish"],
                    [KMGOCRLanguageCodeKey: "nl", KMGOCRLanguageStringKey: "Dutch"],
                    [KMGOCRLanguageCodeKey: "en", KMGOCRLanguageStringKey: "English"],
                    [KMGOCRLanguageCodeKey: "eo", KMGOCRLanguageStringKey: "Esperanto"],
                    [KMGOCRLanguageCodeKey: "et", KMGOCRLanguageStringKey: "Estonian"],
                    [KMGOCRLanguageCodeKey: "tl", KMGOCRLanguageStringKey: "Filipino"],
                    [KMGOCRLanguageCodeKey: "fi", KMGOCRLanguageStringKey: "Finnish"],
                    [KMGOCRLanguageCodeKey: "fr", KMGOCRLanguageStringKey: "French"],
                    [KMGOCRLanguageCodeKey: "gl", KMGOCRLanguageStringKey: "Galician"],
                    [KMGOCRLanguageCodeKey: "ka", KMGOCRLanguageStringKey: "Georgian"],
                    [KMGOCRLanguageCodeKey: "de", KMGOCRLanguageStringKey: "German"],
                    [KMGOCRLanguageCodeKey: "el", KMGOCRLanguageStringKey: "Greek"],
                    [KMGOCRLanguageCodeKey: "gu", KMGOCRLanguageStringKey: "Gujarati"],
                    [KMGOCRLanguageCodeKey: "ht", KMGOCRLanguageStringKey: "Haitian Creole"],
                    [KMGOCRLanguageCodeKey: "ha", KMGOCRLanguageStringKey: "Hausa"],
                    [KMGOCRLanguageCodeKey: "iw", KMGOCRLanguageStringKey: "Hebrew"],
                    [KMGOCRLanguageCodeKey: "hi", KMGOCRLanguageStringKey: "Hindi"],
                    [KMGOCRLanguageCodeKey: "hmn", KMGOCRLanguageStringKey: "Hmong"],
                    [KMGOCRLanguageCodeKey: "hu", KMGOCRLanguageStringKey: "Hungarian"],
                    [KMGOCRLanguageCodeKey: "is", KMGOCRLanguageStringKey: "Icelandic"],
                    [KMGOCRLanguageCodeKey: "ig", KMGOCRLanguageStringKey: "Igbo"],
                    [KMGOCRLanguageCodeKey: "id", KMGOCRLanguageStringKey: "Indonesian"],
                    [KMGOCRLanguageCodeKey: "ga", KMGOCRLanguageStringKey: "Irish"],
                    [KMGOCRLanguageCodeKey: "it", KMGOCRLanguageStringKey: "Italian"],
                    [KMGOCRLanguageCodeKey: "ja", KMGOCRLanguageStringKey: "Japanese"],
                    [KMGOCRLanguageCodeKey: "jw", KMGOCRLanguageStringKey: "Javanese"],
                    [KMGOCRLanguageCodeKey: "kn", KMGOCRLanguageStringKey: "Kannada"],
                    [KMGOCRLanguageCodeKey: "kk", KMGOCRLanguageStringKey: "Kazakh"],
                    [KMGOCRLanguageCodeKey: "km", KMGOCRLanguageStringKey: "Khmer"],
                    [KMGOCRLanguageCodeKey: "ko", KMGOCRLanguageStringKey: "Korean"],
                    [KMGOCRLanguageCodeKey: "lo", KMGOCRLanguageStringKey: "Lao"],
                    [KMGOCRLanguageCodeKey: "la", KMGOCRLanguageStringKey: "Latin"],
                    [KMGOCRLanguageCodeKey: "lv", KMGOCRLanguageStringKey: "Latvian"],
                    [KMGOCRLanguageCodeKey: "lt", KMGOCRLanguageStringKey: "Lithuanian"],
                    [KMGOCRLanguageCodeKey: "mk", KMGOCRLanguageStringKey: "Macedonian"],
                    [KMGOCRLanguageCodeKey: "mg", KMGOCRLanguageStringKey: "Malagasy"],
                    [KMGOCRLanguageCodeKey: "ms", KMGOCRLanguageStringKey: "Malay"],
                    [KMGOCRLanguageCodeKey: "ml", KMGOCRLanguageStringKey: "Malayalam"],
                    [KMGOCRLanguageCodeKey: "mt", KMGOCRLanguageStringKey: "Maltese"],
                    [KMGOCRLanguageCodeKey: "mi", KMGOCRLanguageStringKey: "Maori"],
                    [KMGOCRLanguageCodeKey: "mr", KMGOCRLanguageStringKey: "Marathi"],
                    [KMGOCRLanguageCodeKey: "mn", KMGOCRLanguageStringKey: "Mongolian"],
                    [KMGOCRLanguageCodeKey: "my", KMGOCRLanguageStringKey: "Myanmar (Burmese)"],
                    [KMGOCRLanguageCodeKey: "ne", KMGOCRLanguageStringKey: "Nepali"],
                    [KMGOCRLanguageCodeKey: "no", KMGOCRLanguageStringKey: "Norwegian"],
                    [KMGOCRLanguageCodeKey: "fa", KMGOCRLanguageStringKey: "Persian"],
                    [KMGOCRLanguageCodeKey: "pl", KMGOCRLanguageStringKey: "Polish"],
                    [KMGOCRLanguageCodeKey: "pt", KMGOCRLanguageStringKey: "Portuguese"],
                    [KMGOCRLanguageCodeKey: "ma", KMGOCRLanguageStringKey: "Punjabi"],
                    [KMGOCRLanguageCodeKey: "ro", KMGOCRLanguageStringKey: "Romanian"],
                    [KMGOCRLanguageCodeKey: "ru", KMGOCRLanguageStringKey: "Russian"],
                    [KMGOCRLanguageCodeKey: "sr", KMGOCRLanguageStringKey: "Serbian"],
                    [KMGOCRLanguageCodeKey: "st", KMGOCRLanguageStringKey: "Sesotho"],
                    [KMGOCRLanguageCodeKey: "si", KMGOCRLanguageStringKey: "Sinhala"],
                    [KMGOCRLanguageCodeKey:"sk", KMGOCRLanguageStringKey:"Slovak"],
                    [KMGOCRLanguageCodeKey:"sl", KMGOCRLanguageStringKey:"Slovenian"],
                    [KMGOCRLanguageCodeKey:"so", KMGOCRLanguageStringKey:"Somali"],
                    [KMGOCRLanguageCodeKey:"es", KMGOCRLanguageStringKey:"Spanish"],
                    [KMGOCRLanguageCodeKey:"su", KMGOCRLanguageStringKey:"Sudanese"],
                    [KMGOCRLanguageCodeKey:"sw", KMGOCRLanguageStringKey:"Swahili"],
                    [KMGOCRLanguageCodeKey:"sv", KMGOCRLanguageStringKey:"Swedish"],
                    [KMGOCRLanguageCodeKey:"tg", KMGOCRLanguageStringKey:"Tajik"],
                    [KMGOCRLanguageCodeKey:"ta", KMGOCRLanguageStringKey:"Tamil"],
                    [KMGOCRLanguageCodeKey:"te", KMGOCRLanguageStringKey:"Telugu"],
                    [KMGOCRLanguageCodeKey:"th", KMGOCRLanguageStringKey:"Thai"],
                    [KMGOCRLanguageCodeKey:"tr", KMGOCRLanguageStringKey:"Turkish"],
                    [KMGOCRLanguageCodeKey:"uk", KMGOCRLanguageStringKey:"Ukrainian"],
                    [KMGOCRLanguageCodeKey:"ur", KMGOCRLanguageStringKey:"Urdu"],
                    [KMGOCRLanguageCodeKey:"uz", KMGOCRLanguageStringKey:"Uzbek"],
                    [KMGOCRLanguageCodeKey:"vi", KMGOCRLanguageStringKey:"Vietnamese"],
                    [KMGOCRLanguageCodeKey:"cy", KMGOCRLanguageStringKey:"Welsh"],
                    [KMGOCRLanguageCodeKey:"yi", KMGOCRLanguageStringKey:"Yiddish"],
                    [KMGOCRLanguageCodeKey:"yo", KMGOCRLanguageStringKey:"Yoruba"],
                    [KMGOCRLanguageCodeKey:"zu", KMGOCRLanguageStringKey:"Zulu"]]
        }
        return [[KMGOCRLanguageCodeKey: "en-US", KMGOCRLanguageStringKey: "English"],
                [KMGOCRLanguageCodeKey: "fr-FR", KMGOCRLanguageStringKey: "French"],
                [KMGOCRLanguageCodeKey: "it-IT", KMGOCRLanguageStringKey: "Italian"],
                [KMGOCRLanguageCodeKey: "de-DE", KMGOCRLanguageStringKey: "German"],
                [KMGOCRLanguageCodeKey: "es-ES", KMGOCRLanguageStringKey: "Spanish"],
                [KMGOCRLanguageCodeKey: "pt-BR", KMGOCRLanguageStringKey: "Portuguese"],
                [KMGOCRLanguageCodeKey: "zh-Hant", KMGOCRLanguageStringKey: "Chinese Traditional"],
                [KMGOCRLanguageCodeKey: "zh-Hans", KMGOCRLanguageStringKey: "Chinese Simplified"]
               ]
    }
    
    func recognitionImages(_ images: [Any], withLanguages languages: [Any]) {
        recognitionImages(images, withLanguages: languages, fileType: nil, filePath: nil)
    }
    func recognitionImages(_ images: [Any], withLanguages languages: [Any], fileType: String?, filePath: URL?) {
        self.ocrPath = []
        self.finishIndex = 0
        self.fileType = "PDF"
        self.images = images
        if filePath == nil {
            self.filePath = URL(string: NSSearchPathForDirectoriesInDomains(.desktopDirectory, .userDomainMask, true)[0])
        } else {
            self.filePath = filePath
        }
        
        if self.OCRType == .Google {
            gocrRecognitionImages(images, withLanguages: languages)
        } else {
            if #available(macOS 10.15, *) {
                if appleRequest != nil {
                    appleRequest?.cancel()
                    appleRequest = nil
                }
                recognitionAppleImage(at: self.finishIndex)
            } else {
                self.delegate?.GOCRManagerDidFail?(self, atIndex: self.finishIndex, error: nil)
            }
        }
    }
    func gocrRecognitionImages(_ images: [Any], withLanguages languages: [Any]) {
        if images.isEmpty || images.count == 0 { return }
        
        self.delegate?.GOCRManagerDidStartOCR?(self)
        
        for i in 0..<images.count {
            self.delegate?.GOCRManagerDidStart?(self, atIndex: i)
            let queue = KMOCROperationQueue.sharedInstance
            queue.maxConcurrentOperationCount = 1
            var image = NSImage()
            if images[i] is NSImage {
                image = images[i] as! NSImage
            } else {
                let data = images[i] as! Data
                image = NSImage(data: data) ?? NSImage()
            }
            let op = KMGOCROperation(recognitionImg: image, imgIndex: i)
            op.selectedLanguages = NSMutableArray(array: languages) as? Array<Any>
            op.operationDelegate = self
            queue.addOperation(op)
        }
    }
    func recognitionAppleImage(at index: Int) { 
        guard index < self.images?.count ?? 0 else {
            if #available(macOS 10.15, *) {
                appleRequest?.cancel()
            }
            appleRequest = nil 
            self.delegate?.GOCRManagerDidFinishOCR?(self)
            return
        }
        finishIndex = index
        delegate?.GOCRManagerDidStart?(self, atIndex: self.finishIndex)
        
        if let image = self.images?[finishIndex] as? NSImage {
            recognitionAppleImage(self.images?[self.finishIndex] as! NSImage)
        } else if let data = self.images?[finishIndex] as? Data {
            let image = NSImage(data: data)
            recognitionAppleImage(image!)
        }
    }
    func recognitionAppleImage(_ image: NSImage) {
        DispatchQueue.global().async { [weak self] in guard let self = self else { return }
            self.appleRequest = VNRecognizeTextRequest { [weak self] request, error in
                guard let self = self else { return }
            var results: [KMGOCRResult]? = nil
            if let reqResults = request.results as? [VNRecognizedTextObservation] {
                results = self.responseDataRequest(request, dictionary: nil, imageSize: image.size)
            }
            
            var resultArray = [[String: Any]]()
            if let results = results {
                for result in results {
                    let dic: [String: Any] = ["x": result.textBounds.origin.x,
                                              "y": result.textBounds.origin.y,
                                              "width": result.textBounds.size.width,
                                              "height": result.textBounds.size.height,
                                              "text": result.text]
                    resultArray.append(dic)
                }
            }
            
            DispatchQueue.main.async {
                if let error = error {
                    self.delegate?.GOCRManagerDidFail?(self, atIndex: self.finishIndex, error: error)
                } else {
                    self.delegate?.GOCRManagerDidFinish?(self, atIndex: self.finishIndex, results: results!)
                }
                self.recognitionAppleImage(at: self.finishIndex + 1)
            }
            }
            self.appleRequest?.usesCPUOnly = true
            self.appleRequest?.recognitionLevel = self.appleRecognitionMode ?? .accurate
            
            var array: Array<String> = self.languages ?? []
            if array.contains("zh-Hant") {
                array.removeAll(where: { $0 == "zh-Hant" })
                array.insert("zh-Hant", at: 0)
            }
            if array.contains("zh-Hans") {
                array.removeAll(where: { $0 == "zh-Hans" })
                array.insert("zh-Hans", at: 0)
            }
            if array.isEmpty {
                array = ["zh-Hans", "zh-Hant"]
            }
            self.appleRequest?.recognitionLanguages = array
            
            let options = [VNImageOption: Any]()
            if let cgImage = self.nsImageToCGImageRef(image) {
                let handler = VNImageRequestHandler(cgImage: cgImage, options: options)
                try? handler.perform([self.appleRequest!])
            }
        }
    }
    
    func nsImageToCGImageRef(_ image: NSImage) -> CGImage? {
        guard let imageData = image.tiffRepresentation else { return nil }
        if let imageSource = CGImageSourceCreateWithData(imageData as CFData, nil), let imageRef = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) {
            return imageRef
        }
        return nil
    }
    
    func responseDataRequest(_ request: VNRequest, dictionary: NSDictionary?, imageSize: CGSize) -> [KMGOCRResult] {
        var results: Array<KMGOCRResult> = Array()
        let maximumCandidates = 1
        var OCRStr = ""
//        if let observations = request.results as? [VNRecognizedTextObservation] {
//            for observation in observations {
//                if let text = observation.topCandidates(maximumCandidates).first {
//                    OCRStr.append("\(text.string)\n")
//                    
//                    var x: CGFloat = 0, y: CGFloat = 0, width: CGFloat = 0, height: CGFloat = 0
//                    var error: Error?
//                    let cnt = text.string.count
//                    let range = 0 ..< cnt
//
//                    if let rectangleObservation = text.boundingBox(for: Range(location: 0, length: text.string.count), error: &error) {
//                        x = rectangleObservation.topLeft.x * imageSize.width
//                        y = (1 - rectangleObservation.topLeft.y) * imageSize.height
//                        width = rectangleObservation.boundingBox.size.width * imageSize.width
//                        height = rectangleObservation.boundingBox.size.height * imageSize.height
//                    }
//                    
//                    let result = KMGOCRResult()
//                    result.text = text.string
//                    result.locale = ""
//                    result.textBounds = CGRect(x: x, y: y, width: width, height: height)
//                    results.append(result)
//                }
//            }
//            // Following Google's logic, the first element of the array represents the text of the entire image
//            if results.count > 0 {
//                let result = KMGOCRResult()
//                result.text = OCRStr
//                if self.languages.count > 0 {
//                    result.locale = self.languages[0]
//                }
//                result.textBounds = CGRect.zero
//                results.insert(result, at: 0)
//            }
//        }
        return results
    }
*/
}