Shelfnet works, also visualization. Postprocessing need to be parallelized
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
@@ -117,6 +117,9 @@ target_link_libraries(map_demo tkDNN)
|
||||
add_executable(demo demo/demo/demo.cpp)
|
||||
target_link_libraries(demo tkDNN)
|
||||
|
||||
add_executable(seg_demo demo/demo/seg_demo.cpp)
|
||||
target_link_libraries(seg_demo tkDNN)
|
||||
|
||||
#-------------------------------------------------------------------------------
|
||||
# Install
|
||||
#-------------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
#include <iostream>
|
||||
#include <signal.h>
|
||||
#include <stdlib.h> /* srand, rand */
|
||||
#include <unistd.h>
|
||||
#include <mutex>
|
||||
|
||||
#include "SegmentationNN.h"
|
||||
|
||||
bool gRun;
|
||||
bool SAVE_RESULT = false;
|
||||
|
||||
void sig_handler(int signo) {
|
||||
std::cout<<"request gateway stop\n";
|
||||
gRun = false;
|
||||
}
|
||||
|
||||
int main(int argc, char *argv[]) {
|
||||
|
||||
std::cout<<"detection\n";
|
||||
signal(SIGINT, sig_handler);
|
||||
|
||||
|
||||
std::string net = "shelfnet_fp32.rt";
|
||||
if(argc > 1)
|
||||
net = argv[1];
|
||||
std::string input = "../../ShelfNet/ShelfNet18_realtime/data/leftImg8bit/test/modena/000302.png";
|
||||
if(argc > 2)
|
||||
input = argv[2];
|
||||
int n_batch = 1;
|
||||
if(argc > 3)
|
||||
n_batch = atoi(argv[3]);
|
||||
bool show = false;
|
||||
if(argc > 4)
|
||||
show = atoi(argv[4]);
|
||||
|
||||
if(n_batch < 1 || n_batch > 64)
|
||||
FatalError("Batch dim not supported");
|
||||
|
||||
if(!show)
|
||||
SAVE_RESULT = true;
|
||||
|
||||
int n_classes = 19;
|
||||
|
||||
tk::dnn::SegmentationNN segNN;
|
||||
segNN.init(net, n_classes, n_batch);
|
||||
|
||||
gRun = true;
|
||||
|
||||
cv::VideoCapture cap(input);
|
||||
if(!cap.isOpened())
|
||||
gRun = false;
|
||||
else
|
||||
std::cout<<"camera started\n";
|
||||
|
||||
cv::VideoWriter resultVideo;
|
||||
if(SAVE_RESULT) {
|
||||
int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
|
||||
int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
|
||||
resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(w, h));
|
||||
}
|
||||
|
||||
cv::Mat frame;
|
||||
if(show)
|
||||
cv::namedWindow("segmentation", cv::WINDOW_NORMAL);
|
||||
|
||||
std::vector<cv::Mat> batch_frame;
|
||||
std::vector<cv::Mat> batch_dnn_input;
|
||||
|
||||
while(gRun) {
|
||||
batch_dnn_input.clear();
|
||||
batch_frame.clear();
|
||||
|
||||
for(int bi=0; bi< n_batch; ++bi){
|
||||
cap >> frame;
|
||||
if(!frame.data)
|
||||
break;
|
||||
|
||||
batch_frame.push_back(frame);
|
||||
|
||||
// this will be resized to the net format
|
||||
batch_dnn_input.push_back(frame.clone());
|
||||
}
|
||||
if(!frame.data)
|
||||
break;
|
||||
|
||||
//inference
|
||||
segNN.update(batch_dnn_input, n_batch);
|
||||
segNN.draw();
|
||||
|
||||
if(show){
|
||||
for(int bi=0; bi< n_batch; ++bi){
|
||||
cv::imshow("segmentation", batch_frame[bi]);
|
||||
cv::waitKey(1);
|
||||
}
|
||||
}
|
||||
if(n_batch == 1 && SAVE_RESULT)
|
||||
resultVideo << frame;
|
||||
}
|
||||
|
||||
std::cout<<"segmentation end\n";
|
||||
double mean = 0;
|
||||
|
||||
std::cout<<COL_GREENB<<"\n\nTime stats:\n";
|
||||
std::cout<<"Min: "<<*std::min_element(segNN.stats.begin(), segNN.stats.end())/n_batch<<" ms\n";
|
||||
std::cout<<"Max: "<<*std::max_element(segNN.stats.begin(), segNN.stats.end())/n_batch<<" ms\n";
|
||||
for(int i=0; i<segNN.stats.size(); i++) mean += segNN.stats[i]; mean /= segNN.stats.size();
|
||||
std::cout<<"Avg: "<<mean/n_batch<<" ms\t"<<1000/(mean/n_batch)<<" FPS\n"<<COL_END;
|
||||
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -430,18 +430,22 @@ public:
|
||||
|
||||
};
|
||||
|
||||
enum ResizeMode_t { NEAREST= 0,
|
||||
LINEAR= 1};
|
||||
|
||||
/**
|
||||
Resize layer
|
||||
*/
|
||||
class Resize : public Layer {
|
||||
|
||||
public:
|
||||
Resize(Network *net, int scale_c, int scale_h, int scale_w, bool fixed=false);
|
||||
Resize(Network *net, int scale_c, int scale_h, int scale_w, bool fixed=false, ResizeMode_t mode=NEAREST);
|
||||
virtual ~Resize();
|
||||
virtual layerType_t getLayerType() { return LAYER_RESIZE; };
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
|
||||
ResizeMode_t mode;
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -5,8 +5,8 @@
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map);
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim);
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map, double min=0, double max=0);
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, double min=0, double max=0);
|
||||
cv::Mat vizLayer2Mat(tk::dnn::Network *net, int layer, int imgdim = 1000);
|
||||
|
||||
}}
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
#ifndef SEGMENTATIONNN_H
|
||||
#define SEGMENTATIONNN_H
|
||||
|
||||
#include <iostream>
|
||||
#include <signal.h>
|
||||
#include <stdlib.h>
|
||||
#include <unistd.h>
|
||||
#include <mutex>
|
||||
#include "utils.h"
|
||||
|
||||
#include <opencv2/core/core.hpp>
|
||||
#include <opencv2/highgui/highgui.hpp>
|
||||
#include <opencv2/imgproc/imgproc.hpp>
|
||||
#include <opencv2/core/hal/interface.h>
|
||||
|
||||
|
||||
|
||||
#include "tkdnn.h"
|
||||
#include "NetworkViz.h"
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
class SegmentationNN {
|
||||
|
||||
protected:
|
||||
tk::dnn::NetworkRT *netRT = nullptr;
|
||||
int nBatches = 1;
|
||||
|
||||
std::vector<cv::Size> originalSize;
|
||||
std::vector<cv::Mat> masks;
|
||||
cv::Mat bgr[3];
|
||||
dnnType *input;
|
||||
dnnType *input_d;
|
||||
float* confidences_h;
|
||||
|
||||
/**
|
||||
* This method preprocess the image, before feeding it to the NN.
|
||||
*
|
||||
* @param frame original frame to adapt for inference.
|
||||
* @param bi batch index
|
||||
*/
|
||||
void preprocess(cv::Mat &frame, const int bi=0) {
|
||||
|
||||
frame.convertTo(frame, CV_32FC3, 1 / 255.0, 0);
|
||||
|
||||
cv::split(frame, bgr);
|
||||
float mean[] = {0.485, 0.456, 0.406};
|
||||
float stddev[] = {0.229, 0.224, 0.225};
|
||||
for(int i=0; i<3; i++){
|
||||
bgr[2-i] -= mean[i];
|
||||
bgr[2-i] /= stddev[i];
|
||||
}
|
||||
cv::merge(bgr, 3, frame);
|
||||
|
||||
int crop_size = netRT->input_dim.w;
|
||||
int H = frame.rows;
|
||||
int W = frame.cols;
|
||||
cv::Mat frame_cropped;
|
||||
|
||||
cv::Mat mask(frame.size(), CV_8UC3, cv::Scalar(255,255,255));
|
||||
|
||||
if(H != W){
|
||||
if(H < W){
|
||||
int top = (W - H)/2;
|
||||
int bottom = W - top - H;
|
||||
cv::copyMakeBorder(frame, frame_cropped, top, bottom, 0, 0, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
cv::copyMakeBorder(mask, mask, top, bottom, 0, 0, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
}
|
||||
else{
|
||||
int left = (H - W)/2;
|
||||
int right = H - left - W;
|
||||
cv::copyMakeBorder(frame, frame_cropped, 0, 0, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
cv::copyMakeBorder(mask, mask, 0, 0, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
}
|
||||
}
|
||||
|
||||
resize(frame_cropped, frame_cropped, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
||||
resize(mask, mask, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
||||
masks[bi] = mask.clone();
|
||||
|
||||
cv::split(frame_cropped, bgr);
|
||||
for (int i = 0; i < netRT->input_dim.c; i++){
|
||||
int idx = i * frame_cropped.rows * frame_cropped.cols;
|
||||
int ch = netRT->input_dim.c-1 -i;
|
||||
memcpy((void *)&input[idx + netRT->input_dim.tot()*bi], (void *)bgr[ch].data, frame_cropped.rows * frame_cropped.cols * sizeof(dnnType));
|
||||
}
|
||||
checkCuda(cudaMemcpyAsync(input_d+ netRT->input_dim.tot()*bi, input + netRT->input_dim.tot()*bi, netRT->input_dim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
|
||||
}
|
||||
|
||||
/**
|
||||
* This method postprocess the output of the NN to obtain the correct
|
||||
* boundig boxes.
|
||||
*
|
||||
* @param bi batch index
|
||||
*/
|
||||
void postprocess(const int bi=0) {
|
||||
dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
|
||||
|
||||
dataDim_t odim = netRT->output_dim;
|
||||
|
||||
|
||||
checkCuda(cudaMemcpy(confidences_h, rt_out, odim.tot() * sizeof(float), cudaMemcpyDeviceToHost));
|
||||
|
||||
for(int i=0;i<odim.h;++i){
|
||||
for(int j=0;j<odim.w;++j){
|
||||
float max_conf = 0;
|
||||
int max_id = 0;
|
||||
|
||||
for(int k=0; k<odim.c;++k){
|
||||
float cur_conf = confidences_h[bi*odim.tot()+k*odim.h*odim.w+i*odim.h+j];
|
||||
if(cur_conf > max_conf){
|
||||
max_conf = cur_conf;
|
||||
max_id = k;
|
||||
}
|
||||
}
|
||||
confidences_h[bi*odim.tot()+0*odim.h*odim.w+i*odim.h+j] = max_id;
|
||||
}
|
||||
}
|
||||
dataDim_t vdim = odim;
|
||||
vdim.c = 1;
|
||||
segmented[bi] = vizData2Mat(confidences_h, vdim, 1024, 0, 18);
|
||||
};
|
||||
|
||||
public:
|
||||
int classes = 0;
|
||||
std::vector<double> stats; /*keeps track of inference times (ms)*/
|
||||
std::vector<std::string> classesNames;
|
||||
std::vector<cv::Mat> segmented;
|
||||
|
||||
SegmentationNN() {};
|
||||
~SegmentationNN(){};
|
||||
|
||||
/**
|
||||
* Method used to inialize the class, allocate memory and compute
|
||||
* needed data.
|
||||
*
|
||||
* @param tensor_path path to the rt file og the NN.
|
||||
* @param n_classes number of classes for the given dataset.
|
||||
* @param n_batches maximum number of batches to use in inference
|
||||
* @return true if everything is correct, false otherwise.
|
||||
*/
|
||||
bool init(const std::string& tensor_path, const int n_classes=19, const int n_batches=1){
|
||||
std::cout<<(tensor_path).c_str()<<"\n";
|
||||
if(!fileExist(tensor_path.c_str()))
|
||||
FatalError("This file do not exists" + tensor_path );
|
||||
|
||||
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
|
||||
classes = n_classes;
|
||||
nBatches = n_batches;
|
||||
|
||||
checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
|
||||
checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
|
||||
|
||||
confidences_h = (float *)malloc(netRT->output_dim.tot() * sizeof(float));
|
||||
|
||||
segmented.resize(nBatches);
|
||||
masks.resize(nBatches);
|
||||
}
|
||||
|
||||
/**
|
||||
* This method performs the whole detection of the NN.
|
||||
*
|
||||
* @param frames frames to run detection on.
|
||||
* @param cur_batches number of batches to use in inference
|
||||
* @param save_times if set to true, preprocess, inference and postprocess times
|
||||
* are saved on a csv file, otherwise not.
|
||||
* @param times pointer to the output stream where to write times
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation
|
||||
*/
|
||||
void update(std::vector<cv::Mat>& frames, const int cur_batches=1){
|
||||
if(cur_batches > nBatches)
|
||||
FatalError("A batch size greater than nBatches cannot be used");
|
||||
|
||||
originalSize.clear();
|
||||
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi){
|
||||
if(!frames[bi].data)
|
||||
FatalError("No image data feed to detection");
|
||||
originalSize.push_back(frames[bi].size());
|
||||
preprocess(frames[bi], bi);
|
||||
}
|
||||
TKDNN_TSTOP
|
||||
}
|
||||
|
||||
//do inference
|
||||
tk::dnn::dataDim_t dim = netRT->input_dim;
|
||||
dim.n = cur_batches;
|
||||
{
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
TKDNN_TSTART
|
||||
netRT->infer(dim, input_d);
|
||||
TKDNN_TSTOP
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
stats.push_back(t_ns);
|
||||
}
|
||||
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi)
|
||||
postprocess(bi);
|
||||
TKDNN_TSTOP
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Method to draw boundixg boxes and labels on a frame.
|
||||
*/
|
||||
void draw(const int cur_batches=1) {
|
||||
for(int i=0; i<cur_batches; ++i){
|
||||
|
||||
cv::bitwise_and(segmented[i], masks[i], segmented[i]);
|
||||
|
||||
cv::imshow("segmented", segmented[i]);
|
||||
cv::waitKey(1);
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}}
|
||||
|
||||
#endif /* SEGMENTATIONNN_H*/
|
||||
@@ -36,7 +36,7 @@
|
||||
#define COL_PURPLEB "\033[1;35m"
|
||||
#define COL_CYANB "\033[1;36m"
|
||||
|
||||
#define TKDNN_VERBOSE 0
|
||||
#define TKDNN_VERBOSE 1
|
||||
|
||||
// Simple Timer
|
||||
#define TKDNN_TSTART timespec start, end; \
|
||||
|
||||
@@ -72,6 +72,7 @@ do
|
||||
./test_imuodom &>> $out_file
|
||||
print_output $? imuodom
|
||||
|
||||
test_net shelfnet
|
||||
test_net yolo4
|
||||
test_net yolo4_berkeley
|
||||
test_net yolo3
|
||||
|
||||
+2
-1
@@ -483,6 +483,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Resize *l) {
|
||||
IResizeLayer *lRT = networkRT->addResize(*input); //default is kNEAREST
|
||||
checkNULL(lRT);
|
||||
Dims d{};
|
||||
lRT->setResizeMode(ResizeMode(l->mode));
|
||||
lRT->setOutputDimensions(DimsCHW{l->output_dim.c, l->output_dim.h, l->output_dim.w});
|
||||
return lRT;
|
||||
}
|
||||
@@ -514,7 +515,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
|
||||
|
||||
ITensor *back_tens = tensors[l->backLayer];
|
||||
|
||||
if(false) //l->backLayer->output_dim.c == l->output_dim.c && !l->mul) FIXME
|
||||
if(l->backLayer->output_dim.c == l->output_dim.c && !l->mul)
|
||||
{
|
||||
IElementWiseLayer *lRT = networkRT->addElementWise(*input, *back_tens, ElementWiseOperation::kSUM);
|
||||
checkNULL(lRT);
|
||||
|
||||
+7
-8
@@ -6,23 +6,22 @@
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map) {
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map,double min, double max) {
|
||||
|
||||
if(min == 0 && max == 0)
|
||||
cv::minMaxIdx(map, &min, &max);
|
||||
|
||||
double min;
|
||||
double max;
|
||||
cv::minMaxIdx(map, &min, &max);
|
||||
cv::Mat adjMap;
|
||||
// expand your range to 0..255. Similar to histEq();
|
||||
map.convertTo(adjMap,CV_8UC1, 255 / (max-min), -min);
|
||||
//return adjMap;
|
||||
|
||||
|
||||
cv::Mat falseColorsMap;
|
||||
applyColorMap(adjMap, falseColorsMap, cv::COLORMAP_HOT);
|
||||
applyColorMap(adjMap, falseColorsMap, cv::COLORMAP_VIRIDIS);
|
||||
return falseColorsMap;
|
||||
}
|
||||
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim) {
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim, double min, double max) {
|
||||
dnnType *data = nullptr;
|
||||
|
||||
// copy to CPU
|
||||
@@ -38,7 +37,7 @@ cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim) {
|
||||
cv::Mat grid = cv::Mat(gridSize, CV_8UC3, cv::Scalar(0));
|
||||
|
||||
for(int i=0; i<dim.c;i++) {
|
||||
cv::Mat raw = vizFloat2colorMap(cv::Mat(cv::Size(dim.w, dim.h),CV_32FC1, data + dim.w*dim.h*i));
|
||||
cv::Mat raw = vizFloat2colorMap(cv::Mat(cv::Size(dim.w, dim.h),CV_32FC1, data + dim.w*dim.h*i), min, max);
|
||||
int r = i / gridDim;
|
||||
int c = i - r * gridDim;
|
||||
raw.copyTo(grid.rowRange(r*dim.h, r*dim.h + dim.h).colRange(c*dim.w, c*dim.w + dim.w));
|
||||
|
||||
+2
-1
@@ -5,8 +5,9 @@
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
Resize::Resize(Network *net, int scale_c, int scale_h, int scale_w, bool fixed) : Layer(net) {
|
||||
Resize::Resize(Network *net, int scale_c, int scale_h, int scale_w, bool fixed, ResizeMode_t mode) : Layer(net) {
|
||||
|
||||
this->mode = mode;
|
||||
if(fixed){
|
||||
output_dim.c = scale_c;
|
||||
output_dim.h = scale_h;
|
||||
|
||||
@@ -191,7 +191,7 @@ int main()
|
||||
down_out.push_back(l_last);
|
||||
|
||||
new tk::dnn::Conv2d (&net, out_channel*2, 3, 3, 2, 2, 1, 1, ladder[li++], false);
|
||||
last = new tk::dnn::Activation (&net, CUDNN_ACTIVATION_RELU);
|
||||
last = new tk::dnn::Activation (&net, tk::dnn::ACTIVATION_LEAKY, 0.0f, 0.0f); //should be ReLU
|
||||
}
|
||||
|
||||
new tk::dnn::Conv2d (&net, 256, 3, 3, 1, 1, 1, 1, ladder[li++], true, false, 1, true);
|
||||
@@ -231,12 +231,12 @@ int main()
|
||||
new tk::dnn::Conv2d (&net, 64, 3, 3, 1, 1, 1, 1, conv_out[ci++], true);
|
||||
new tk::dnn::Activation (&net, tk::dnn::ACTIVATION_LEAKY, 0.0f, 0.01);
|
||||
new tk::dnn::Conv2d (&net, 19, 3, 3, 1, 1, 1, 1, conv_out[ci++], false);
|
||||
// /*up_out[i] =*/ new tk::dnn::Resize(&net, 19, net.input_dim.h, net.input_dim.w, true);
|
||||
/*up_out[i] =*/ new tk::dnn::Resize(&net, 19, net.input_dim.h, net.input_dim.w, true, tk::dnn::ResizeMode_t::LINEAR);
|
||||
// }
|
||||
|
||||
// new tk::dnn::Softmax(&net);
|
||||
new tk::dnn::Softmax(&net);
|
||||
|
||||
const char *output_bin = "shelfnet/debug/conv_out-conv_out.bin";
|
||||
const char *output_bin = "shelfnet/debug/softmax.bin";
|
||||
|
||||
// Load input
|
||||
dnnType *data;
|
||||
|
||||
@@ -31,7 +31,7 @@ int main(int argc, char *argv[]) {
|
||||
std::cout<<"Testing with batchsize: "<<BATCH_SIZE<<"\n";
|
||||
printCenteredTitle(" TENSORRT inference ", '=', 30);
|
||||
float total_time = 0;
|
||||
for(int i=0; i<1200; i++) {
|
||||
for(int i=0; i<64; i++) {
|
||||
|
||||
// generate input
|
||||
for(int j=0; j<netRT.input_dim.tot(); j++) {
|
||||
@@ -58,6 +58,6 @@ int main(int argc, char *argv[]) {
|
||||
}
|
||||
}
|
||||
}
|
||||
std::cout<<"avg: "<<total_time/1200.<<std::endl;
|
||||
std::cout<<"avg: "<<total_time/64.<<std::endl;
|
||||
return ret_tensorrt;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user