File: MNISTDataHandler.mm

package info (click to toggle)
onnxruntime 1.23.2%2Bdfsg-3
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 340,744 kB
  • sloc: cpp: 3,222,135; python: 188,267; ansic: 114,318; asm: 37,927; cs: 36,849; java: 10,962; javascript: 6,811; pascal: 4,126; sh: 2,996; xml: 705; objc: 281; makefile: 67
file content (164 lines) | stat: -rw-r--r-- 5,657 bytes parent folder | download | duplicates (4)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#import "MNISTDataHandler.h"
#import "OnnxruntimeModule.h"
#import "TensorHelper.h"
#import <Foundation/Foundation.h>
#import <React/RCTLog.h>

NS_ASSUME_NONNULL_BEGIN

@implementation MNISTDataHandler

RCT_EXPORT_MODULE(MNISTDataHandler)

// It returns mode path in local device,
// so that onnxruntime is able to load a model using a given path.
RCT_EXPORT_METHOD(getLocalModelPath : (RCTPromiseResolveBlock)resolve rejecter : (RCTPromiseRejectBlock)reject) {
  @try {
    NSString* modelPath = [[NSBundle mainBundle] pathForResource:@"mnist" ofType:@"ort"];
    NSFileManager* fileManager = [NSFileManager defaultManager];
    if ([fileManager fileExistsAtPath:modelPath]) {
      resolve(modelPath);
    } else {
      reject(@"mnist", @"no such a model", nil);
    }
  } @catch (NSException* exception) {
    reject(@"mnist", @"no such a model", nil);
  }
}

// It returns image path.
RCT_EXPORT_METHOD(getImagePath : (RCTPromiseResolveBlock)resolve reject : (RCTPromiseRejectBlock)reject) {
  @try {
    NSString* imagePath = [[NSBundle mainBundle] pathForResource:@"3" ofType:@"jpg"];
    NSFileManager* fileManager = [NSFileManager defaultManager];
    if ([fileManager fileExistsAtPath:imagePath]) {
      resolve(imagePath);
    } else {
      reject(@"mnist", @"no such an image", nil);
    }
  } @catch (NSException* exception) {
    reject(@"mnist", @"no such an image", nil);
  }
}

// It gets raw input data, which can be uri or byte array and others,
// returns cooked data formatted as input of a model.
RCT_EXPORT_METHOD(preprocess : (NSString*)uri resolve : (RCTPromiseResolveBlock)resolve reject : (RCTPromiseRejectBlock)reject) {
  @try {
    NSDictionary* inputDataMap = [self preprocess:uri];
    resolve(inputDataMap);
  } @catch (NSException* exception) {
    reject(@"mnist", @"can't load an image", nil);
  }
}

// It gets a result from onnxruntime and a duration of session time for input data,
// returns output data formatted as React Native map.
RCT_EXPORT_METHOD(postprocess : (NSDictionary*)result resolve : (RCTPromiseResolveBlock)resolve reject : (RCTPromiseRejectBlock)reject) {
  @try {
    NSDictionary* cookedMap = [self postprocess:result];
    resolve(cookedMap);
  } @catch (NSException* exception) {
    reject(@"mnist", @"can't pose-process an image", nil);
  }
}

- (NSDictionary*)preprocess:(NSString*)uri {
  UIImage* image = [UIImage imageNamed:@"3.jpg"];

  CGSize scale = CGSizeMake(28, 28);
  UIGraphicsBeginImageContextWithOptions(scale, NO, 1.0);
  [image drawInRect:CGRectMake(0, 0, scale.width, scale.height)];
  UIImage* scaledImage = UIGraphicsGetImageFromCurrentImageContext();
  UIGraphicsEndImageContext();

  CGImageRef imageRef = [scaledImage CGImage];
  NSUInteger width = CGImageGetWidth(imageRef);
  NSUInteger height = CGImageGetHeight(imageRef);
  CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();

  const NSUInteger rawDataSize = height * width * 4;
  std::vector<unsigned char> rawData(rawDataSize);
  NSUInteger bytesPerPixel = 4;
  NSUInteger bytesPerRow = bytesPerPixel * width;
  CGContextRef context = CGBitmapContextCreate(rawData.data(), width, height, 8, bytesPerRow, colorSpace,
                                               kCGImageAlphaPremultipliedLast | kCGImageByteOrder32Big);
  CGColorSpaceRelease(colorSpace);
  CGContextSetBlendMode(context, kCGBlendModeCopy);
  CGContextDrawImage(context, CGRectMake(0, 0, width, height), imageRef);
  CGContextRelease(context);

  const NSInteger dimSize = height * width;
  const NSInteger byteBufferSize = dimSize * sizeof(float);

  unsigned char* byteBuffer = static_cast<unsigned char*>(malloc(byteBufferSize));
  NSData* byteBufferRef = [NSData dataWithBytesNoCopy:byteBuffer length:byteBufferSize];
  float* floatPtr = (float*)[byteBufferRef bytes];
  for (NSUInteger h = 0; h < height; ++h) {
    for (NSUInteger w = 0; w < width; ++w) {
      NSUInteger byteIndex = (bytesPerRow * h) + w * bytesPerPixel;
      *floatPtr++ = rawData[byteIndex];
    }
  }
  floatPtr = (float*)[byteBufferRef bytes];

  NSMutableDictionary* inputDataMap = [NSMutableDictionary dictionary];

  NSMutableDictionary* inputTensorMap = [NSMutableDictionary dictionary];

  // dims
  NSArray* dims = @[
    [NSNumber numberWithInt:1],
    [NSNumber numberWithInt:1],
    [NSNumber numberWithInt:static_cast<int>(height)],
    [NSNumber numberWithInt:static_cast<int>(width)]
  ];
  inputTensorMap[@"dims"] = dims;

  // type
  inputTensorMap[@"type"] = JsTensorTypeFloat;

  // encoded data
  NSString* data = [byteBufferRef base64EncodedStringWithOptions:0];
  inputTensorMap[@"data"] = data;

  inputDataMap[@"Input3"] = inputTensorMap;

  return inputDataMap;
}

- (NSDictionary*)postprocess:(NSDictionary*)result {
  NSMutableString* detectionResult = [NSMutableString string];

  NSDictionary* outputTensor = [result objectForKey:@"Plus214_Output_0"];

  NSString* data = [outputTensor objectForKey:@"data"];
  NSData* buffer = [[NSData alloc] initWithBase64EncodedString:data options:0];
  float* values = (float*)[buffer bytes];
  int count = (int)[buffer length] / 4;

  int argmax = 0;
  float maxValue = 0.0f;
  for (int i = 0; i < count; ++i) {
    if (values[i] > maxValue) {
      maxValue = values[i];
      argmax = i;
    }
  }

  if (maxValue == 0.0f) {
    detectionResult = [NSMutableString stringWithString:@"No match"];
  } else {
    detectionResult = [NSMutableString stringWithFormat:@"%d", argmax];
  }

  NSDictionary* cookedMap = @{@"result" : detectionResult};
  return cookedMap;
}

@end

NS_ASSUME_NONNULL_END