PDF_Office/PDF Office/PDF Master/Class/PDFTools/OCR/Manager/KMGOCROperation.m

//
//  KMGOCROperation.m
//  PDF Reader Pro Edition
//
//  Created by 万军 on 2020/3/23.
//

#import "KMGOCROperation.h"
#import "KMGOCRManager.h"
#import <AppKit/AppKit.h>

#define KMGOC_API_URL     @"https://vision.googleapis.com/v1/images:annotate"

#if VERSION_FREE
#define KMGOC_API_KEY     @"AIzaSyBhSRohpngAu8pSgFDXPytslNDHgGm7uDs"
#else
#define KMGOC_API_KEY     @"AIzaSyCJuqJ9YvtkFKMl1mW3Yq-av3mmI9ScbRY"
#endif

@interface KMGOCROperation ()

@property (assign, nonatomic, getter = isExecuting)     BOOL executing;
@property (assign, nonatomic, getter = isFinished)      BOOL finished;
@property (assign, nonatomic, getter = isCancelled)     BOOL cancelled;

@property (nonatomic, strong) NSString *fileName;
@property (nonatomic, strong) NSImage *orcImage;
@property (nonatomic, strong) NSURLSessionDataTask *task;
@property (nonatomic, assign) NSInteger imageIndex;

@end

@implementation KMGOCROperation
- (void)dealloc {

}

- (instancetype)initWithRecognitionImage:(NSImage *)image withImageIndex:(NSInteger)imageIndex {
    self = [super init];
    if (self) {

        self.fileName = [self fileNameWithDate];
        self.imageIndex = imageIndex;
        self.orcImage = image;
        
        self.queuePriority = NSOperationQueuePriorityNormal;
        
        self.name = self.fileName;
        self.executing = NO;
        self.finished = NO;
    }
    return self;
}

#pragma mark - overwrite
- (void)start {
    if ([self p_checkCancelled]) {
        return;
    }
    
    self.executing  = YES;
    
    [NSThread detachNewThreadSelector:@selector(main) toTarget:self withObject:nil];
}

- (void)main {
    @try {
        if ([self p_checkCancelled]) {
            return;
        }
        [self recognitionImage:_orcImage];
        
        while (self.executing) {
            if ([self p_checkCancelled]) {
                return;
            }
        }
    }
    @catch (NSException * e) {
        NSLog(@"Exception %@", e);
    }
}

- (void)cancel
{
    [super cancel];
    
    if (self.task) {
        [self.task cancel];
        self.task = nil;
    }
    
    if ([self.operationDelegate respondsToSelector:@selector(GOCROperation:cancelOCRImageAtIndex:)]) {
        [self.operationDelegate GOCROperation:self cancelOCRImageAtIndex:self.imageIndex];
    }
    
    if (self.executing) {
        self.executing = NO;
        self.finished = YES;
    } else {
        self.finished = NO;
    }
    self.cancelled = YES;
}

#pragma mark - private methods
- (void)p_done
{
    self.executing  = NO;
    self.finished   = YES;
}

- (BOOL)p_checkCancelled
{
    if (self.cancelled) {
        self.finished = YES;
        return YES;
    }
    return NO;
}

- (void)recognitionImage:(NSImage *)image {
    if ([self.operationDelegate respondsToSelector:@selector(GOCROperation:startOCRImageAtIndex:)]) {
        [self.operationDelegate GOCROperation:self startOCRImageAtIndex:self.imageIndex];
    }
    NSString *binaryImageData = [self base64EncodeImage:image];
    NSLog(@"binaryImageData == %@",binaryImageData);
    if (!binaryImageData) {
        if ([self.operationDelegate respondsToSelector:@selector(GOCROperation:failureOCRImageAtIndex:error:)]) {
            [self.operationDelegate GOCROperation:self failureOCRImageAtIndex:self.imageIndex error:nil];
        }
        return;
    }
    
    NSString *urlString = [NSString stringWithFormat:@"%@?key=%@",KMGOC_API_URL,KMGOC_API_KEY];
    NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:[NSURL URLWithString:urlString]];
    [request setHTTPMethod: @"POST"];
    [request addValue:@"application/json" forHTTPHeaderField:@"Content-Type"];
    
    NSDictionary *imageDictionary = @{@"content":binaryImageData};
    NSArray *featuresArray = @[@{@"type":@"TEXT_DETECTION",@"maxResults":@10}];
    NSDictionary *paramsDictionary = @{@"requests":@[@{@"image":imageDictionary,
                                                       @"features":featuresArray}]};
    if (self.selectedLanguages && self.selectedLanguages.count>0) {
        NSDictionary *imageContextDictionary = @{@"languageHints":self.selectedLanguages};
        paramsDictionary = @{@"requests":@[@{@"image":imageDictionary,
                                             @"features":featuresArray,
                                             @"imageContext":imageContextDictionary}]};
    }
    
    NSError *error;
    NSData *requestData = [NSJSONSerialization dataWithJSONObject:paramsDictionary options:0 error:&error];
    [request setHTTPBody:requestData];
    
    NSURLSession *URLSession = [NSURLSession sharedSession];
    self.task = [URLSession dataTaskWithRequest:request
                              completionHandler:^ (NSData *data, NSURLResponse *response, NSError *error) {
        if (NSURLErrorCancelled == error.code) {
          return ;
        }
        NSArray *results = nil;
        if (!error) {
          NSDictionary *dictionary = [NSJSONSerialization JSONObjectWithData:data options:NSJSONReadingMutableLeaves error:nil];
          results = [self responseDataResults:dictionary];
        }

        dispatch_async(dispatch_get_main_queue(), ^{
          if (error || !results) {
              if ([self.operationDelegate respondsToSelector:@selector(GOCROperation:failureOCRImageAtIndex:error:)]) {
                  [self.operationDelegate GOCROperation:self failureOCRImageAtIndex:self.imageIndex error:error];
              }
          } else {
              if ([self.operationDelegate respondsToSelector:@selector(GOCROperation:finishOCRImageAtIndex:results:)]) {
                  [self.operationDelegate GOCROperation:self finishOCRImageAtIndex:self.imageIndex results:results];
              }
          }
        });
        if (error || !results) {
            [self cancel];
        } else {
            [self p_done];
        }
    }];
    [self.task resume];
}

#pragma mark - Private Methods
- (NSString *)fileNameWithDate
{
    NSDateFormatter *formatter = [[NSDateFormatter alloc ] init];
    [formatter setDateFormat:@"YYYY-MM-dd-hh-mm-ss-SSS"];
    NSString *dateString = [formatter stringFromDate:[NSDate date]];
    NSString *fileName = [NSString stringWithFormat:@"%@ %ld",dateString, (long)_imageIndex];
    return fileName;
}

- (NSData *)compressImageData:(NSData *)imageData toMaxFileSize:(NSInteger)maxFileSize {
    CGFloat compression = 0.9f;
    CGFloat maxCompression = 0.1f;
    NSData *compressImageData = imageData;
    while ([imageData length] > maxFileSize && compression > maxCompression) {
        compression -= 0.1;
        NSBitmapImageRep *imageRep = [NSBitmapImageRep imageRepWithData:imageData];
        compressImageData = [imageRep representationUsingType:NSBitmapImageFileTypeJPEG
                                                   properties:@{NSImageCompressionFactor:@(compression)}];
    }
    return compressImageData;
}

- (NSString *)base64EncodeImage:(NSImage *)image {
    NSData *data = [image TIFFRepresentation];
    NSBitmapImageRep *imageRep = [NSBitmapImageRep imageRepWithData:data];
    [imageRep setSize:[image size]];
    NSData *imagedata = [imageRep representationUsingType:NSBitmapImageFileTypePNG properties:@{}];
    
    // Resize the image if it exceeds the 4MB API limit
    if ([imagedata length] > 4194304) {
        imagedata = [self compressImageData:imagedata toMaxFileSize:4194304];
    }
    
    NSString *base64String = nil;
    if ([imagedata respondsToSelector:@selector(base64EncodedStringWithOptions:)]) {
        base64String = [imagedata base64EncodedStringWithOptions:NSDataBase64EncodingEndLineWithCarriageReturn];
    } else {
        base64String = [imagedata base64Encoding];
    }
    return base64String;
}

- (NSArray *)responseDataResults:(NSDictionary *)dictionary {
    NSArray *responses = [dictionary isKindOfClass:[NSDictionary class]] ? [dictionary objectForKey:@"responses"] : nil;
    NSDictionary *responseData = [responses isKindOfClass:[NSArray class]] ? [responses firstObject] : nil;
    NSDictionary *errorObj = [dictionary isKindOfClass:[NSDictionary class]] ? [dictionary objectForKey:@"error"] : nil;
    
    if (errorObj) {
        return nil;
    }
    
    NSMutableArray *results  = nil;
    NSArray *textAnnotations = [responseData isKindOfClass:[NSDictionary class]] ? [responseData objectForKey:@"textAnnotations"] : nil;
    if (textAnnotations && [textAnnotations isKindOfClass:[NSArray class]]) {
        results = [NSMutableArray array];
        for (NSDictionary *annotation in textAnnotations) {
            CGRect textBounds = CGRectZero;
            NSArray *vertices = annotation[@"boundingPoly"][@"vertices"];
            if (vertices && [vertices isKindOfClass:[NSArray class]]) {
                CGFloat minX = 0, minY = 0, maxX = 0, maxY = 0;
                for (int i=0; i<vertices.count; i++) {
                    NSDictionary *vertex = vertices[i];
                    CGFloat x = [vertex[@"x"] floatValue];
                    CGFloat y = [vertex[@"y"] floatValue];
                    minX = i == 0 ? x : MIN(x, minX);
                    minY = i == 0 ? y : MIN(y, minY);
                    maxX = i == 0 ? x : MAX(x, maxX);
                    maxY = i == 0 ? y : MAX(y, maxY);
                }
                textBounds = CGRectMake(minX, minY, maxX-minX, maxY-minY);
            }
            
            KMGOCRResult *result = [[KMGOCRResult alloc] init];
            result.text = annotation[@"description"];
            result.locale = annotation[@"locale"];
            result.textBounds = textBounds;
            [results addObject:result];
        }
    }
    return results;
}
@end