Skip to content

Commit

Permalink
Create thread for mask calculation
Browse files Browse the repository at this point in the history
  • Loading branch information
@peckto
peckto committed May 30, 2021
1 parent a3e1b0e commit ad80c97
Showing 1 changed file with 119 additions and 39 deletions.
158 changes: 119 additions & 39 deletions deepseg.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <string>
#include <thread>
#include <mutex>
#include <condition_variable>

#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
Expand Down Expand Up @@ -145,16 +146,19 @@ typedef std::chrono::high_resolution_clock::time_point timestamp_t;
typedef struct {
timestamp_t bootns;
timestamp_t lastns;
timestamp_t waitns;
timestamp_t lockns;
timestamp_t copyns;
timestamp_t prepns;
timestamp_t tfltns;
timestamp_t maskns;
timestamp_t postns;
timestamp_t v4l2ns;
timestamp_t grabns;
timestamp_t retrns;
// values form ai thread are already converted to ns
long ai_condns;
long ai_prens;
long ai_tfltns;
long ai_postns;
long ai_maskns;
long ai_totalns;
} timinginfo_t;

timestamp_t timestamp() {
Expand Down Expand Up @@ -191,10 +195,19 @@ typedef struct {
cv::Rect roidim;
cv::Mat mask;
cv::Mat mroi;
cv::Mat raw;
cv::Mat raw1;
cv::Mat raw2;
cv::Mat *raw_current;
cv::Mat *raw_next;
cv::Mat ofinal;
cv::Size blur;
float ratio;
timinginfo_t *pti;
std::mutex lock_raw;
std::mutex lock_mask;
std::condition_variable condition_new_frame;
bool new_frame;
bool running;
} calcinfo_t;

modeltype_t get_modeltype(const char* modelname) {
Expand Down Expand Up @@ -262,11 +275,37 @@ void init_tensorflow(calcinfo_t &info) {

// create Mat for small mask
info.ofinal = cv::Mat(info.output.rows,info.output.cols,CV_8UC1);

info.new_frame = false;
info.running = false;

info.raw_next = &info.raw1;
info.raw_current = &info.raw2;
}

void calc_mask(calcinfo_t &info, timinginfo_t &ti) {
void calc_mask(calcinfo_t &info) {
timestamp_t t = timestamp();
cv::Mat *raw_tmp;

// wait for new frame to arrive
timestamp_t t0 = timestamp();
{
std::unique_lock<std::mutex> hold(info.lock_raw);
while (!info.new_frame) {
info.condition_new_frame.wait(hold);
}
info.pti->ai_condns = diffnanosecs(timestamp(), t0);
t0 = timestamp();

// change frame buffer pointer
info.new_frame = false;
raw_tmp = info.raw_next;
info.raw_next = info.raw_current;
info.raw_current = raw_tmp;
}

// map ROI
cv::Mat roi = info.raw(info.roidim);
cv::Mat roi = (*info.raw_current)(info.roidim);

// resize ROI to input size
cv::Mat in_u8_bgr, in_u8_rgb;
Expand All @@ -283,11 +322,14 @@ void calc_mask(calcinfo_t &info, timinginfo_t &ti) {

// convert to float and normalize to values expected by model
in_u8_rgb.convertTo(info.input,CV_32FC3,info.norm.scaling,info.norm.offset);
ti.prepns=timestamp();
info.pti->ai_prens = diffnanosecs(timestamp(), t0);
t0 = timestamp();


// Run inference
TFLITE_MINIMAL_CHECK(interpreter->Invoke() == kTfLiteOk);
ti.tfltns=timestamp();
info.pti->ai_tfltns = diffnanosecs(timestamp(), t0);
t0 = timestamp();

float* tmp = (float*)info.output.data;
uint8_t* out = (uint8_t*)info.ofinal.data;
Expand Down Expand Up @@ -337,20 +379,36 @@ void calc_mask(calcinfo_t &info, timinginfo_t &ti) {
fprintf(stderr, "Unknown model type\n");
break;
}
ti.maskns=timestamp();
info.pti->ai_postns = diffnanosecs(timestamp(), t0);
t0 = timestamp();

// scale up into full-sized mask
cv::Mat tmpbuf;
cv::resize(info.ofinal,tmpbuf,cv::Size(info.raw.rows/info.ratio,info.raw.rows));
{
cv::Mat tmpbuf;
std::lock_guard<std::mutex> hold(info.lock_mask);
cv::resize(info.ofinal,tmpbuf,cv::Size(info.raw_current->rows/info.ratio,info.raw_current->rows));
// blur at full size for maximum smoothness
cv::blur(tmpbuf,info.mroi,info.blur);
}

// blur at full size for maximum smoothness
cv::blur(tmpbuf,info.mroi,info.blur);
info.pti->ai_maskns = diffnanosecs(timestamp(), t0);
info.pti->ai_totalns = diffnanosecs(timestamp(), t);
}

bool is_number(const std::string &s) {
return !s.empty() && std::all_of(s.begin(), s.end(), ::isdigit);
}

void *calc_mask_thread(void *args) {

calcinfo_t *calcinfo = (calcinfo_t *)args;
while (!calcinfo->running)
usleep(10000);
while (calcinfo->running)
calc_mask(*calcinfo);
return NULL;
}

int main(int argc, char* argv[]) {

printf("deepseg version %s\n", _STR(DEEPSEG_VERSION));
Expand Down Expand Up @@ -532,6 +590,7 @@ int main(int argc, char* argv[]) {
if (fourcc)
cap.set(CV_CAP_PROP_FOURCC, fourcc);
cap.set(CV_CAP_PROP_CONVERT_RGB, true);
cap.set(CV_CAP_PROP_BUFFERSIZE, 2);

auto modeltype = get_modeltype(modelname);
auto norm = get_normalization(modeltype);
Expand All @@ -540,12 +599,18 @@ int main(int argc, char* argv[]) {
exit(1);
}
calcinfo_t calcinfo = { modelname, modeltype, norm, threads, width, height, debug };
calcinfo.pti = &ti;
init_tensorflow(calcinfo);

cv::Mat mask;
cv::Mat raw;
ti.lastns = timestamp();
printf("Startup: %ldns\n", diffnanosecs(ti.lastns,ti.bootns));

std::thread ai(calc_mask_thread, &calcinfo);

bool filterActive = true;
calcinfo.running = true;

// mainloop
for(bool running = true; running; ) {
Expand All @@ -556,43 +621,51 @@ int main(int argc, char* argv[]) {
ti.grabns = timestamp();

// copy new frame to buffer
cap.retrieve(calcinfo.raw);
cap.retrieve(raw);
ti.retrns = timestamp();
ti.copyns = timestamp();
ti.lockns=timestamp();

if (calcinfo.raw.rows == 0 || calcinfo.raw.cols == 0) continue; // sanity check
if (raw.rows == 0 || raw.cols == 0) continue; // sanity check

// copy frame to ai thread local buffer
{
std::lock_guard<std::mutex> hold(calcinfo.lock_raw);
ti.lockns=timestamp();
*calcinfo.raw_next = raw.clone();
calcinfo.new_frame = true;
calcinfo.condition_new_frame.notify_all();
}

ti.copyns=timestamp();

if (blur_strength) {
calcinfo.raw.copyTo(bg);
raw.copyTo(bg);
cv::GaussianBlur(bg,bg,cv::Size(blur_strength,blur_strength),0);
}

if (filterActive) {
// do background detection magic
calc_mask(calcinfo, ti);
{
std::lock_guard<std::mutex> hold(calcinfo.lock_mask);
mask = calcinfo.mask.clone();
}

// alpha blend background over foreground using mask
calcinfo.raw = alpha_blend(bg, calcinfo.raw, calcinfo.mask);
} else {
// fix up timing values
ti.maskns=ti.tfltns=ti.prepns=ti.copyns;
raw = alpha_blend(bg, raw, mask);
}

if (flipHorizontal && flipVertical) {
cv::flip(calcinfo.raw,calcinfo.raw,-1);
cv::flip(raw,raw,-1);
} else if (flipHorizontal) {
cv::flip(calcinfo.raw,calcinfo.raw,1);
cv::flip(raw,raw,1);
} else if (flipVertical) {
cv::flip(calcinfo.raw,calcinfo.raw,0);
cv::flip(raw,raw,0);
}
ti.postns=timestamp();

// write frame to v4l2loopback as YUYV
calcinfo.raw = convert_rgb_to_yuyv(calcinfo.raw);
int framesize = calcinfo.raw.step[0]*calcinfo.raw.rows;
raw = convert_rgb_to_yuyv(raw);
int framesize = raw.step[0]*raw.rows;
while (framesize > 0) {
int ret = write(lbfd,calcinfo.raw.data,framesize);
int ret = write(lbfd,raw.data,framesize);
TFLITE_MINIMAL_CHECK(ret > 0);
framesize -= ret;
}
Expand All @@ -607,23 +680,28 @@ int main(int argc, char* argv[]) {
}

// timing details..
printf("lock:%9ld [grab:%9ld retr:%9ld] copy:%9ld open:%9ld tflt:%9ld mask:%9ld post:%9ld v4l2:%9ld FPS: %5.2f\e[K\r",
diffnanosecs(ti.lockns,ti.retrns),

printf("video [grab:%9ld retr:%9ld lock:%9ld copy:%9ld post:%9ld v4l2:%9ld] ai [cond:%9ld pre:%9ld tflt:%9ld post:%9ld mask:%9ld] - FPS: %5.2f (video) %5.2f (ai)\e[K\r",
diffnanosecs(ti.grabns, ti.lastns),
diffnanosecs(ti.retrns,ti.grabns),
diffnanosecs(ti.lockns,ti.retrns),
diffnanosecs(ti.copyns,ti.lockns),
diffnanosecs(ti.prepns,ti.copyns),
diffnanosecs(ti.tfltns,ti.prepns),
diffnanosecs(ti.maskns,ti.tfltns),
diffnanosecs(ti.postns,ti.maskns),
diffnanosecs(ti.postns,ti.copyns),
diffnanosecs(ti.v4l2ns,ti.postns),
1e9/diffnanosecs(ti.v4l2ns,ti.lastns));
ti.ai_condns,
ti.ai_prens,
ti.ai_tfltns,
ti.ai_postns,
ti.ai_maskns,
1e9/diffnanosecs(ti.v4l2ns,ti.lastns),
1000000000./ti.ai_totalns
);
fflush(stdout);
ti.lastns = timestamp();
if (debug < 2) continue;

cv::Mat test;
cv::cvtColor(calcinfo.raw,test,CV_YUV2BGR_YUYV);
cv::cvtColor(raw,test,CV_YUV2BGR_YUYV);
cv::imshow("DeepSeg " _STR(DEEPSEG_VERSION),test);

auto keyPress = cv::waitKey(1);
Expand All @@ -644,5 +722,7 @@ int main(int argc, char* argv[]) {
}

printf("\n");
calcinfo.running = false;
ai.join();
return 0;
}

0 comments on commit ad80c97

Please sign in to comment.