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 Apr 4, 2021
1 parent 094b355 commit 3499d60
Showing 1 changed file with 78 additions and 27 deletions.
105 changes: 78 additions & 27 deletions deepseg.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,16 +118,16 @@ 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 openns;
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 openns;
long tfltns;
long maskns;
} timinginfo_t;

timestamp_t timestamp() {
Expand Down Expand Up @@ -161,6 +161,10 @@ typedef struct {
cv::Mat ofinal;
cv::Mat element;
float ratio;
timinginfo_t *pti;
pthread_mutex_t lock_raw;
pthread_mutex_t lock_mask;
bool start;
} calcinfo_t;

void init_tensorflow(calcinfo_t &info) {
Expand Down Expand Up @@ -198,11 +202,23 @@ void init_tensorflow(calcinfo_t &info) {

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

info.lock_raw = PTHREAD_MUTEX_INITIALIZER;
info.lock_mask = PTHREAD_MUTEX_INITIALIZER;
info.start = false;
}

void calc_mask(calcinfo_t &info, timinginfo_t &ti) {
void calc_mask(calcinfo_t &info) {
timestamp_t t0 = timestamp();

// copy frame from shared buffer
cv::Mat raw;
pthread_mutex_lock(&info.lock_raw);
raw = info.raw.clone();
pthread_mutex_unlock(&info.lock_raw);

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

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

// convert to float and normalize values to [-1;1]
in_u8_rgb.convertTo(info.input,CV_32FC3,1.0/128.0,-1.0);
ti.openns=timestamp();
info.pti->openns = diffnanosecs(timestamp(), t0);
t0 = timestamp();


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

float* tmp = (float*)info.output.data;
uint8_t* out = (uint8_t*)info.ofinal.data;
Expand Down Expand Up @@ -271,15 +289,26 @@ void calc_mask(calcinfo_t &info, timinginfo_t &ti) {
out[n] = (val & 0xE0) | (out[n] >> 3);
}
}
ti.maskns=timestamp();
info.pti->maskns = diffnanosecs(timestamp(), t0);

// denoise
cv::Mat tmpbuf;
cv::dilate(info.ofinal,tmpbuf,info.element);
cv::erode(tmpbuf,info.ofinal,info.element);

// scale up into full-sized mask
cv::resize(info.ofinal,info.mroi,cv::Size(info.raw.rows/info.ratio,info.raw.rows));
pthread_mutex_lock(&info.lock_mask);
cv::resize(info.ofinal,info.mroi,cv::Size(raw.rows/info.ratio,raw.rows));
pthread_mutex_unlock(&info.lock_mask);
}

void *calc_mask_thread(void *args) {
calcinfo_t *calcinfo = (calcinfo_t *)args;
while (!calcinfo->start) usleep(10000);
while (calcinfo->start) {
calc_mask(*calcinfo);
}
return NULL;
}

int main(int argc, char* argv[]) {
Expand Down Expand Up @@ -437,11 +466,22 @@ int main(int argc, char* argv[]) {
cap.set(CV_CAP_PROP_CONVERT_RGB, true);

calcinfo_t calcinfo = { modelname, threads, width, height, debug };
calcinfo.pti = &ti;
init_tensorflow(calcinfo);

// kick off separate grabber thread to keep OpenCV/FFMpeg happy (or it lags badly)
cv::Mat mask;
cv::Mat raw;
ti.lastns = timestamp();
printf("Startup: %ldns\n", diffnanosecs(ti.lastns,ti.bootns));

pthread_t ai;
if (pthread_create(&ai, NULL, calc_mask_thread, &calcinfo)) {
perror("creating calc_mask thread");
exit(1);
}


bool filterActive = true;

// mainloop
Expand All @@ -453,35 +493,44 @@ 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();

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

// copy frame to ai thread local buffer
pthread_mutex_lock(&calcinfo.lock_raw);
ti.lockns=timestamp();
calcinfo.raw = raw.clone();
pthread_mutex_unlock(&calcinfo.lock_raw);
ti.copyns=timestamp();

if (calcinfo.raw.rows == 0 || calcinfo.raw.cols == 0) continue; // sanity check
// signal ai thread data arrived for the first time
calcinfo.start = true;

if (filterActive) {
// do background detection magic
calc_mask(calcinfo, ti);
pthread_mutex_lock(&calcinfo.lock_mask);
mask = calcinfo.mask.clone();
pthread_mutex_unlock(&calcinfo.lock_mask);

// copy background over raw cam image using mask
bg.copyTo(calcinfo.raw,calcinfo.mask);
bg.copyTo(raw,calcinfo.mask);
} // filterActive

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 @@ -501,10 +550,10 @@ int main(int argc, char* argv[]) {
diffnanosecs(ti.grabns, ti.lastns),
diffnanosecs(ti.retrns,ti.grabns),
diffnanosecs(ti.copyns,ti.lockns),
diffnanosecs(ti.openns,ti.copyns),
diffnanosecs(ti.tfltns,ti.openns),
diffnanosecs(ti.maskns,ti.tfltns),
diffnanosecs(ti.postns,ti.maskns),
ti.openns,
ti.tfltns,
ti.maskns,
diffnanosecs(ti.postns,ti.copyns),
diffnanosecs(ti.v4l2ns,ti.postns));

int e2 = cv::getTickCount();
Expand All @@ -515,7 +564,7 @@ int main(int argc, char* argv[]) {
if (debug < 2) continue;

cv::Mat test;
cv::cvtColor(calcinfo.raw,test,CV_YUV2BGR_YUYV);
cv::cvtColor(raw,test,CV_YUV2BGR_YUYV);
cv::imshow("output.png",test);

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

printf("\n");
calcinfo.start = false;
pthread_join(ai, NULL);
return 0;
}

0 comments on commit 3499d60

Please sign in to comment.