mirror of
https://github.com/jomjol/AI-on-the-edge-device.git
synced 2026-01-27 12:50:39 +03:00
412 lines
8.4 KiB
C++
412 lines
8.4 KiB
C++
#include "defines.h"
|
|
|
|
#include "CTfLiteClass.h"
|
|
#include "ClassLogFile.h"
|
|
#include "Helper.h"
|
|
#include "psram.h"
|
|
#include "esp_log.h"
|
|
|
|
#include <sys/stat.h>
|
|
|
|
static const char *TAG = "TFLITE";
|
|
|
|
CTfLiteClass::CTfLiteClass()
|
|
{
|
|
model = nullptr;
|
|
modelfile = NULL;
|
|
interpreter = nullptr;
|
|
input = nullptr;
|
|
output = nullptr;
|
|
kTensorArenaSize = TENSOR_ARENA_SIZE;
|
|
tensor_arena = (uint8_t *)psram_get_shared_tensor_arena_memory();
|
|
}
|
|
|
|
CTfLiteClass::~CTfLiteClass()
|
|
{
|
|
delete interpreter;
|
|
psram_free_shared_tensor_arena_and_model_memory();
|
|
}
|
|
|
|
bool CTfLiteClass::MakeStaticResolver(void)
|
|
{
|
|
if (resolver.AddFullyConnected() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddFullyConnected() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddReshape() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddReshape() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddSoftmax() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddSoftmax() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddConv2D() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddConv2D() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddMaxPool2D() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddMaxPool2D() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddQuantize() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddQuantize() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddMul() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddMul() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddAdd() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddAdd() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddLeakyRelu() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddLeakyRelu() failed");
|
|
return false;
|
|
}
|
|
if (resolver.AddDequantize() != kTfLiteOk)
|
|
{
|
|
ESP_LOGE(TAG, "load AddDequantize() failed");
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
float CTfLiteClass::GetOutputValue(int nr)
|
|
{
|
|
TfLiteTensor *output2 = interpreter->output(0);
|
|
|
|
int numer_output = output2->dims->data[1];
|
|
if ((nr + 1) > numer_output)
|
|
{
|
|
return -1000;
|
|
}
|
|
|
|
return output2->data.f[nr];
|
|
}
|
|
|
|
int CTfLiteClass::GetClassFromImageBasis(CImageBasis *rs)
|
|
{
|
|
if (!LoadInputImageBasis(rs))
|
|
{
|
|
return -1000;
|
|
}
|
|
|
|
Invoke();
|
|
|
|
return GetOutClassification();
|
|
}
|
|
|
|
int CTfLiteClass::GetOutClassification(int _von, int _bis)
|
|
{
|
|
TfLiteTensor *output2 = interpreter->output(0);
|
|
|
|
float zw_max;
|
|
float zw;
|
|
int zw_class;
|
|
|
|
if (output2 == NULL)
|
|
{
|
|
return -1;
|
|
}
|
|
|
|
int numeroutput = output2->dims->data[1];
|
|
// ESP_LOGD(TAG, "number output neurons: %d", numeroutput);
|
|
|
|
if (_bis == -1)
|
|
{
|
|
_bis = numeroutput - 1;
|
|
}
|
|
|
|
if (_von == -1)
|
|
{
|
|
_von = 0;
|
|
}
|
|
|
|
if (_bis >= numeroutput)
|
|
{
|
|
ESP_LOGD(TAG, "NUMBER OF OUTPUT NEURONS does not match required classification!");
|
|
return -1;
|
|
}
|
|
|
|
zw_max = output2->data.f[_von];
|
|
zw_class = _von;
|
|
for (int i = _von + 1; i <= _bis; ++i)
|
|
{
|
|
zw = output2->data.f[i];
|
|
if (zw > zw_max)
|
|
{
|
|
zw_max = zw;
|
|
zw_class = i;
|
|
}
|
|
}
|
|
return (zw_class - _von);
|
|
}
|
|
|
|
void CTfLiteClass::GetInputDimension(bool silent = false)
|
|
{
|
|
TfLiteTensor *input2 = interpreter->input(0);
|
|
|
|
int numdim = input2->dims->size;
|
|
if (!silent)
|
|
{
|
|
ESP_LOGD(TAG, "NumDimension: %d", numdim);
|
|
}
|
|
|
|
int sizeofdim;
|
|
for (int j = 0; j < numdim; ++j)
|
|
{
|
|
sizeofdim = input2->dims->data[j];
|
|
if (!silent)
|
|
{
|
|
ESP_LOGD(TAG, "SizeOfDimension %d: %d", j, sizeofdim);
|
|
}
|
|
if (j == 1)
|
|
{
|
|
im_height = sizeofdim;
|
|
}
|
|
if (j == 2)
|
|
{
|
|
im_width = sizeofdim;
|
|
}
|
|
if (j == 3)
|
|
{
|
|
im_channel = sizeofdim;
|
|
}
|
|
}
|
|
}
|
|
|
|
int CTfLiteClass::ReadInputDimenstion(int _dim)
|
|
{
|
|
if (_dim == 0)
|
|
{
|
|
return im_width;
|
|
}
|
|
if (_dim == 1)
|
|
{
|
|
return im_height;
|
|
}
|
|
if (_dim == 2)
|
|
{
|
|
return im_channel;
|
|
}
|
|
|
|
return -1;
|
|
}
|
|
|
|
int CTfLiteClass::GetAnzOutPut(bool silent)
|
|
{
|
|
TfLiteTensor *output2 = interpreter->output(0);
|
|
|
|
int numdim = output2->dims->size;
|
|
if (!silent)
|
|
{
|
|
ESP_LOGD(TAG, "NumDimension: %d", numdim);
|
|
}
|
|
|
|
int sizeofdim;
|
|
for (int j = 0; j < numdim; ++j)
|
|
{
|
|
sizeofdim = output2->dims->data[j];
|
|
if (!silent)
|
|
{
|
|
ESP_LOGD(TAG, "SizeOfDimension %d: %d", j, sizeofdim);
|
|
}
|
|
}
|
|
|
|
float fo;
|
|
|
|
// Process the inference results.
|
|
int numeroutput = output2->dims->data[1];
|
|
for (int i = 0; i < numeroutput; ++i)
|
|
{
|
|
fo = output2->data.f[i];
|
|
if (!silent)
|
|
{
|
|
ESP_LOGD(TAG, "Result %d: %f", i, fo);
|
|
}
|
|
}
|
|
return numeroutput;
|
|
}
|
|
|
|
void CTfLiteClass::Invoke()
|
|
{
|
|
if (interpreter != nullptr)
|
|
{
|
|
interpreter->Invoke();
|
|
}
|
|
}
|
|
|
|
bool CTfLiteClass::LoadInputImageBasis(CImageBasis *rs)
|
|
{
|
|
unsigned int w = rs->width;
|
|
unsigned int h = rs->height;
|
|
unsigned char red, green, blue;
|
|
// ESP_LOGD(TAG, "Image: %s size: %d x %d\n", _fn.c_str(), w, h);
|
|
|
|
input_i = 0;
|
|
float *input_data_ptr = (interpreter->input(0))->data.f;
|
|
|
|
for (int y = 0; y < h; ++y)
|
|
{
|
|
for (int x = 0; x < w; ++x)
|
|
{
|
|
red = rs->GetPixelColor(x, y, 0);
|
|
green = rs->GetPixelColor(x, y, 1);
|
|
blue = rs->GetPixelColor(x, y, 2);
|
|
|
|
*(input_data_ptr) = (float)red;
|
|
input_data_ptr++;
|
|
|
|
*(input_data_ptr) = (float)green;
|
|
input_data_ptr++;
|
|
|
|
*(input_data_ptr) = (float)blue;
|
|
input_data_ptr++;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool CTfLiteClass::MakeAllocate(void)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::MakeAllocate");
|
|
|
|
if (!MakeStaticResolver())
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - resolver could not be loaded!");
|
|
return false;
|
|
}
|
|
|
|
if (!model)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - no model loaded!");
|
|
return false;
|
|
}
|
|
|
|
if (model->version() != TFLITE_SCHEMA_VERSION)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "The selected model does not match the tflite schema version!");
|
|
return false;
|
|
}
|
|
|
|
if (!tensor_arena)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - tensor_arena not allocate");
|
|
return false;
|
|
}
|
|
|
|
interpreter = new tflite::MicroInterpreter(model, resolver, tensor_arena, kTensorArenaSize);
|
|
// LogFile.WriteToFile(ESP_LOG_INFO, TAG, "Trying to load the model. If it crashes here, it ist most likely due to a corrupted model!");
|
|
|
|
if (interpreter)
|
|
{
|
|
TfLiteStatus allocate_status = interpreter->AllocateTensors();
|
|
if (allocate_status != kTfLiteOk)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "AllocateTensors() failed");
|
|
GetInputDimension();
|
|
return false;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "new tflite::MicroInterpreter failed");
|
|
LogFile.WriteHeapInfo("CTfLiteClass::MakeAllocate-new tflite::MicroInterpreter failed");
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
long CTfLiteClass::GetFileSize(std::string filename)
|
|
{
|
|
struct stat stat_buf;
|
|
long rc = -1;
|
|
|
|
FILE *pFile = fopen(filename.c_str(), "rb"); // previously only "rb
|
|
|
|
if (pFile != NULL)
|
|
{
|
|
rc = stat(filename.c_str(), &stat_buf);
|
|
fclose(pFile);
|
|
}
|
|
|
|
return (rc == 0 ? stat_buf.st_size : -1);
|
|
}
|
|
|
|
bool CTfLiteClass::ReadFileToModel(std::string filename)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::ReadFileToModel: " + filename);
|
|
|
|
long size = GetFileSize(filename);
|
|
if (size == -1)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "Model file doesn't exist: " + filename + "!");
|
|
return false;
|
|
}
|
|
else if (size > MAX_MODEL_SIZE)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "Unable to load model '" + filename + "'! It does not fit in the reserved shared memory in PSRAM!");
|
|
return false;
|
|
}
|
|
|
|
LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "Loading Model " + filename + " /size: " + std::to_string(size) + " bytes...");
|
|
|
|
modelfile = (unsigned char *)psram_get_shared_model_memory();
|
|
if (modelfile != NULL)
|
|
{
|
|
FILE *pFile = fopen(filename.c_str(), "rb"); // previously only "rb
|
|
if (pFile != NULL)
|
|
{
|
|
fread(modelfile, 1, size, pFile);
|
|
fclose(pFile);
|
|
|
|
return true;
|
|
}
|
|
else
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::ReadFileToModel: Model does not exist");
|
|
return false;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::ReadFileToModel: Can't allocate enough memory: " + std::to_string(size));
|
|
LogFile.WriteHeapInfo("CTfLiteClass::ReadFileToModel");
|
|
|
|
return false;
|
|
}
|
|
}
|
|
|
|
bool CTfLiteClass::LoadModel(std::string filename)
|
|
{
|
|
LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::LoadModel");
|
|
|
|
if (!ReadFileToModel(filename.c_str()))
|
|
{
|
|
return false;
|
|
}
|
|
|
|
model = tflite::GetModel(modelfile);
|
|
|
|
if (model == nullptr)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|