Vadim Goroshevsky
Vadim Goroshevsky

Reputation: 1556

How to set OpenCV Mat as Tensorflow Lite input and output?

I'm trying to use GPU Delegate in Tensorflow Lite on iOS. My model has inputs and outputs as OpenCV BGR image ([258, 540, 3]). How can I set inputs and outputs in C++ tensorflow lite interpreter? I tried to use this code

int input = interpreter->inputs()[0];
float* out = interpreter->typed_tensor<float>(input);
NSData* slicedData = [self inputDataFromCvMat:slicedImage];
uint8_t* in = (uint8_t*) slicedData.bytes; 
ProcessInputWithFloatModel(in, out, WIDTH, HEIGHT, CHANNELS);
void ProcessInputWithFloatModel(uint8_t* input, float* buffer, int image_width, int image_height, int image_channels) {
  for (int y = 0; y < wanted_input_height; ++y) {
    float* out_row = buffer + (y * wanted_input_width * wanted_input_channels);
    for (int x = 0; x < wanted_input_width; ++x) {
      const int in_x = (y * image_width) / wanted_input_width;
      const int in_y = (x * image_height) / wanted_input_height;
      uint8_t* input_pixel =
      input + (in_y * image_width * image_channels) + (in_x * image_channels);
      float* out_pixel = out_row + (x * wanted_input_channels);
      for (int c = 0; c < wanted_input_channels; ++c) {
        out_pixel[c] = (input_pixel[c] - input_mean) / input_std;
      }
    }
  }
}
- (NSData *)inputDataFromCvMat:(Mat)image {
  NSMutableData *inputData = [[NSMutableData alloc] initWithCapacity:0];
  for (int row = 0; row < HEIGHT + 10; row++) {
    for (int col = 0; col < WIDTH + 10; col++) {
      Vec3b intensity = image.at<Vec3b>(row, col);
      int blue = intensity.val[0];
      int green = intensity.val[1];
      int red = intensity.val[2];
      // we need to put pixel values in BGR (model was trained with opencv)
      [inputData appendBytes:&blue length:sizeof(blue)];
      [inputData appendBytes:&green length:sizeof(green)];
      [inputData appendBytes:&red length:sizeof(red)];
    }
  }
  return inputData;
}

but I don't know what is wrong

Upvotes: 2

Views: 2473

Answers (1)

Vadim Goroshevsky
Vadim Goroshevsky

Reputation: 1556

After some research, I managed to get it working

const int wanted_input_width = 258;
const int wanted_input_height = 540;
const int wanted_input_channels = 3;

Mat image = ...

// write to input
int input = interpreter->inputs()[0];
float* out = interpreter->typed_tensor<float>(input);
uint8_t* in = image.ptr<uint8_t>(0);
ProcessInputWithFloatModel(in, out);
// run interpreter
if (interpreter->Invoke() != kTfLiteOk) {
   LOG(FATAL) << "Failed to invoke!";
}
// get output      
int output_idx = interpreter->outputs()[0];
float* output = interpreter->typed_output_tensor<float>(output_idx);
Mat outputMat = ProcessOutputWithFloatModel(output);
/// Preprocess the input image and feed the TFLite interpreter buffer for a float model.
void ProcessInputWithFloatModel(uint8_t* input, float* buffer) {
  for (int y = 0; y < wanted_input_height; ++y) {
    float* out_row = buffer + (y * wanted_input_width * wanted_input_channels);
    for (int x = 0; x < wanted_input_width; ++x) {
      uint8_t* input_pixel = input + (y * wanted_input_width * wanted_input_channels) + (x * wanted_input_channels);
      float* out_pixel = out_row + (x * wanted_input_channels);
      for (int c = 0; c < wanted_input_channels; ++c) {
        out_pixel[c] = input_pixel[c] / 255.0f;
      }
    }
  }
}

Mat ProcessOutputWithFloatModel(float* input) {
  cv::Mat image = cv::Mat::zeros(wanted_input_height, wanted_input_width, CV_8UC3);
  for (int y = 0; y < wanted_input_height; ++y) {
    for (int x = 0; x < wanted_input_width; ++x) {
      float* input_pixel = input + (y * wanted_input_width * wanted_input_channels) + (x * wanted_input_channels);
      cv::Vec3b & color = image.at<cv::Vec3b>(cv::Point(x, y));
      color[0] = (uchar) floor(input_pixel[0] * 255.0f);
      color[1] = (uchar) floor(input_pixel[1] * 255.0f);
      color[2] = (uchar) floor(input_pixel[2] * 255.0f);
    }
  }
  return image;
}

Upvotes: 7

Related Questions