Smart-Mask-Enforcement-System.ino

/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <TensorFlowLite.h>

#include "main_functions.h"

#include "detection_responder.h"
#include "image_provider.h"
#include "model_settings.h"
#include "person_detect_model_data.h"

#include "audio_provider.h"
#include "feature_provider.h"
#include "micro_features_micro_model_settings.h"
#include "micro_features_model.h"
#include "recognize_commands.h"

#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/version.h"

#include "Arduino.h"

// Globals, used for compatibility with Arduino-style sketches.
namespace {
tflite::ErrorReporter* error_reporter = nullptr;

const tflite::Model* vww_model = nullptr;
tflite::MicroInterpreter* vww_interpreter = nullptr;
TfLiteTensor* vww_input = nullptr;
TfLiteTensor* vww_output = nullptr;

const tflite::Model* kws_model = nullptr;
tflite::MicroInterpreter* kws_interpreter = nullptr;
TfLiteTensor* kws_input = nullptr;
TfLiteTensor* kws_output = nullptr;
FeatureProvider* feature_provider = nullptr;
RecognizeCommands* recognizer = nullptr;
int32_t previous_time = 0;
int8_t feature_buffer[kFeatureElementCount];
int8_t* model_input_buffer = nullptr;

// An area of memory to use for input, output, and intermediate arrays.
constexpr int kTensorArenaSize = 150 * 1024;
static uint8_t tensor_arena[kTensorArenaSize];
}  // namespace

// The name of this function is important for Arduino compatibility.
void setup() {
  // Set up logging. Google style is to avoid globals or statics because of
  // lifetime uncertainty, but since this has a trivial destructor it's okay.
  // NOLINTNEXTLINE(runtime-global-variables)
  static tflite::MicroErrorReporter micro_error_reporter;
  error_reporter = &micro_error_reporter;

  //create the allocator that will be shared between models
  tflite::MicroAllocator* allocator =
      tflite::MicroAllocator::Create(tensor_arena, kTensorArenaSize,
                                              error_reporter);

  // Map the model into a usable data structure. This doesn't involve any
  // copying or parsing, it's a very lightweight operation.
  vww_model = tflite::GetModel(g_person_detect_model_data);
  if (vww_model->version() != TFLITE_SCHEMA_VERSION) {
    TF_LITE_REPORT_ERROR(error_reporter,
                         "Model provided is schema version %d not equal "
                         "to supported version %d.",
                         vww_model->version(), TFLITE_SCHEMA_VERSION);
    return;
  }

  // Pull in only the operation implementations we need.
  // NOLINTNEXTLINE(runtime-global-variables)
  // static tflite::MicroMutableOpResolver<6> micro_op_resolver;
  // micro_op_resolver.AddAveragePool2D();
  // micro_op_resolver.AddConv2D();
  // micro_op_resolver.AddDepthwiseConv2D();
  // micro_op_resolver.AddReshape();
  // micro_op_resolver.AddSoftmax();
  // micro_op_resolver.AddFullyConnected();

  static tflite::AllOpsResolver micro_op_resolver;

  // Build an interpreter to run the VWW model
  // NOLINTNEXTLINE(runtime-global-variables)
  static tflite::MicroInterpreter vww_static_interpreter(
      vww_model, micro_op_resolver, allocator, error_reporter);
  vww_interpreter = &vww_static_interpreter;

  //allocate the VWW model from tensor_arena
  TfLiteStatus allocate_status = vww_interpreter->AllocateTensors();
  if (allocate_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "AllocateTensors() failed");
    return;
  }
  
  vww_input = vww_interpreter->input(0);
  vww_output = vww_interpreter->output(0);

 //Initialize the kws model
  kws_model = tflite::GetModel(g_model);
  if (kws_model->version() != TFLITE_SCHEMA_VERSION) {
    TF_LITE_REPORT_ERROR(error_reporter,
                         "Model provided is schema version %d not equal "
                         "to supported version %d.",
                         kws_model->version(), TFLITE_SCHEMA_VERSION);
    return;
  }
  

  // Build an interpreter to run the KWS model. The allocator is reused
  // NOLINTNEXTLINE(runtime-global-variables)
  static tflite::MicroInterpreter kws_static_interpreter(
      kws_model, micro_op_resolver, allocator, error_reporter);
  kws_interpreter = &kws_static_interpreter;

  //allocate the KWS model from tensor_arena. Some head space is saved
  allocate_status = kws_interpreter->AllocateTensors();
  if (allocate_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "AllocateTensors() failed");
    return;
  }

  kws_input = kws_interpreter->input(0);
  kws_output = kws_interpreter->output(0);
  model_input_buffer = kws_input->data.int8;


  // Prepare to access the audio spectrograms from a microphone or other source
  // that will provide the inputs to the neural network.
  // NOLINTNEXTLINE(runtime-global-variables)
  static FeatureProvider static_feature_provider(kFeatureElementCount,
                                                 feature_buffer);
  feature_provider = &static_feature_provider;

  static RecognizeCommands static_recognizer(error_reporter);
  recognizer = &static_recognizer;

  previous_time = 0;
}

// The name of this function is important for Arduino compatibility.
void loop() {
  // Fetch the spectrogram for the current time.
  const int32_t current_time = LatestAudioTimestamp();
  int how_many_new_slices = 0;
  TfLiteStatus feature_status = feature_provider->PopulateFeatureData(
      error_reporter, previous_time, current_time, &how_many_new_slices);
  if (feature_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "Feature generation failed");
    return;
  }
  previous_time = current_time;
  // If no new audio samples have been received since last time, don't bother
  // running the network model.
  if (how_many_new_slices == 0) {
    return;
  }

  // Copy feature buffer to input tensor
  for (int i = 0; i < kFeatureElementCount; i++) {
    model_input_buffer[i] = feature_buffer[i];
  }

  // Run the model on the spectrogram input and make sure it succeeds.
  TfLiteStatus invoke_status = kws_interpreter->Invoke();
  if (invoke_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed");
    return;
  }
  

  const char* found_command = nullptr;
  uint8_t score = 0;
  bool is_new_command = false;
  TfLiteStatus process_status = recognizer->ProcessLatestResults(
      kws_output, current_time, &found_command, &score, &is_new_command);
  if (process_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter,
                         "RecognizeCommands::ProcessLatestResults() failed");
    return;
  }

  bool heardYes = RespondToKWS(error_reporter, found_command, is_new_command, score);

  if(heardYes){
    //Our keyword spotting model heard 'yes' so we detect if a person is visible 
    
     // Get image from provider.
    if (kTfLiteOk != GetImage(error_reporter, kNumCols, kNumRows, kNumChannels,
                              vww_input->data.int8)) {
      TF_LITE_REPORT_ERROR(error_reporter, "Image capture failed.");
    }
    else{
      TF_LITE_REPORT_ERROR(error_reporter, "Image captured successfully.");
    }
  
    // Run the model on this input and make sure it succeeds.
    if (kTfLiteOk != vww_interpreter->Invoke()) {
      TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed.");
    }
    else{
      TF_LITE_REPORT_ERROR(error_reporter, "Invoke succeeded.");
    }
  
    // Process the inference results.
    int8_t person_score = vww_output->data.uint8[kMaskedScore];
    int8_t no_person_score = vww_output->data.uint8[kUnmaskedScore];
    RespondToDetection(error_reporter, person_score, no_person_score);
  }
  
}