cailean
3 months ago
8 changed files with 146 additions and 49 deletions
@ -0,0 +1,70 @@ |
|||||
|
#include "ofMain.h" |
||||
|
#include "Onnx.h" |
||||
|
|
||||
|
class ModelThread : public ofThread |
||||
|
{ |
||||
|
public: |
||||
|
ofImage* img; |
||||
|
ofFbo* fbo; |
||||
|
Onnx* model; |
||||
|
|
||||
|
|
||||
|
~ModelThread(){ |
||||
|
stop(); |
||||
|
waitForThread(false); |
||||
|
} |
||||
|
|
||||
|
void setup(ofImage* _img, ofFbo* _fbo, Onnx* _model){ |
||||
|
std::lock_guard<std::mutex> lock(mutex); |
||||
|
this->img = _img; |
||||
|
this->fbo = _fbo; |
||||
|
this->model = _model; |
||||
|
} |
||||
|
|
||||
|
void start(){ |
||||
|
startThread(); |
||||
|
} |
||||
|
|
||||
|
void stop(){ |
||||
|
stopThread(); |
||||
|
condition.notify_all(); |
||||
|
} |
||||
|
|
||||
|
void threadedFunction(){ |
||||
|
while(isThreadRunning()){ |
||||
|
std::unique_lock<std::mutex> lock(mutex); |
||||
|
inferDepthImage(fbo, img, model); |
||||
|
condition.wait(lock); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
void update(){ |
||||
|
std::lock_guard<std::mutex> lock(mutex); |
||||
|
condition.notify_one(); |
||||
|
} |
||||
|
|
||||
|
void inferDepthImage(ofFbo* fbo, ofImage* img, Onnx* model){ |
||||
|
std::cout << "infer" << std::endl; |
||||
|
auto output_tensors = model->Run(*img); |
||||
|
float* output_ptr = output_tensors.front().GetTensorMutableData<float>(); |
||||
|
size_t num_elements = output_tensors.front().GetTensorTypeAndShapeInfo().GetElementCount(); |
||||
|
std::cout << "done" << std::endl; |
||||
|
|
||||
|
float min_value = model->ReduceMin(output_ptr, num_elements); |
||||
|
float max_value = model->ReduceMax(output_ptr, num_elements); |
||||
|
|
||||
|
std::cout << "done1" << std::endl; |
||||
|
|
||||
|
model->Normalize(output_ptr, num_elements, min_value, max_value); |
||||
|
|
||||
|
std::cout << "don2e" << std::endl; |
||||
|
|
||||
|
model->DataToFbo(output_ptr, 518, 518, *fbo); |
||||
|
|
||||
|
std::cout << "done3" << std::endl; |
||||
|
} |
||||
|
|
||||
|
|
||||
|
protected: |
||||
|
std::condition_variable condition; |
||||
|
}; |
Loading…
Reference in new issue