gavin/VideoPipe
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
陈赣
2026-06-03 12:43:14 +08:00
commit ba76cfae28
608 changed files with 120791 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
third_party/trt_vehicle/CMakeLists.txt vendored Normal file
View File

@@ -0,0 +1,72 @@
cmake_minimum_required(VERSION 3.10)
project(trt_vehicle VERSION 1.0)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_BUILD_TYPE Debug)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -fPIC -w -fdiagnostics-color=always -pthread")
# save all libs to the same directory
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/libs)
# TensorRT required
set(TRT_LIB_PATH "/usr/local/tensorRT/lib") # change this line if possible
set(TRT_INC_PATH "/usr/local/tensorRT/include") # change this line if possible
link_directories(${TRT_LIB_PATH})
include_directories(${TRT_INC_PATH})
if(DEFINED VP_BUILD_FROM) # tell videopipe where is trt
set(VP_TRT_LIB_PATH ${TRT_LIB_PATH} PARENT_SCOPE)
set(VP_TRT_INC_PATH ${TRT_INC_PATH} PARENT_SCOPE)
endif()
# CUDA required
find_package(CUDA REQUIRED)
message(STATUS "CUDA library status:")
message(STATUS " version: ${CUDA_VERSION}")
message(STATUS " libraries: ${CUDA_LIBRARIES}")
message(STATUS " include path: ${CUDA_INCLUDE_DIRS}")
include_directories(${CUDA_INCLUDE_DIRS})
# OpenCV required
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV library status:")
message(STATUS " version: ${OpenCV_VERSION}")
message(STATUS " libraries: ${OpenCV_LIBS}")
message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")
include_directories(${OpenCV_INCLUDE_DIRS})
# collect dependent libs for trt_vehicle
set(TRT_VEHICLE_DEPEND_LIBS ${OpenCV_LIBS} ${CUDA_LIBRARIES} nvinfer stdc++fs)
# collect source files for trt_vehicle
file(GLOB_RECURSE MODELS "models/*.cpp")
file(GLOB_RECURSE UTIL "util/*.cpp")
#...#
list(APPEND TRT_VEHICLE_CPPS ${MODELS} ${UTIL})
# build for trt_vehicle
add_library(${PROJECT_NAME} SHARED ${TRT_VEHICLE_CPPS})
target_link_libraries(${PROJECT_NAME} ${TRT_VEHICLE_DEPEND_LIBS})
# build samples for trt_vehicle
if(NOT DEFINED VP_BUILD_FROM)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/samples) # save all exe to 'samples'
add_executable(vehicle_cluster "samples/vehicle_cluster.cpp")
target_link_libraries(vehicle_cluster ${PROJECT_NAME})
add_executable(vehicle_color_type "samples/vehicle_color_type.cpp")
target_link_libraries(vehicle_color_type ${PROJECT_NAME})
add_executable(vehicle_compare "samples/vehicle_compare.cpp")
target_link_libraries(vehicle_compare ${PROJECT_NAME})
add_executable(vehicle_plate "samples/vehicle_plate.cpp")
target_link_libraries(vehicle_plate ${PROJECT_NAME})
add_executable(vehicle_scan "samples/vehicle_scan.cpp")
target_link_libraries(vehicle_scan ${PROJECT_NAME})
add_executable(vehicle_search "samples/vehicle_search.cpp")
target_link_libraries(vehicle_search ${PROJECT_NAME})
add_executable(vehicle "samples/vehicle.cpp")
target_link_libraries(vehicle ${PROJECT_NAME})
endif()

54
third_party/trt_vehicle/README.md vendored Executable file
View File

@@ -0,0 +1,54 @@
## What `trt_vehicle` can do:
1. Vehicle detector
2. Vehicle plate detector and recognizer
3. Vehicle scanner based on side view of body
4. Vehicle color and type classifier
5. Vehicle feature encoder used for search(1:N) or comparison(1:1)
## How to install tensorrt and cuda ?
Refer to NVIDIA official web
1. download files according to your GPUs (below for this repo) `cuda_11.1.0_455.23.05_linux.run`, `cudnn-11.1-linux-x64-v8.0.5.39.tgz`, `TensorRT-7.2.1.6.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.0.tar.gz`
2. run `./cuda_11.1.0_455.23.05_linux.run` install cuda(at `/usr/local`) and driver, maybe need to reboot machine some times
3. upzip `cudnn-11.1-linux-x64-v8.0.5.39.tgz` , copy all header files to '/usr/local/cuda/include' and copy all lib files to '/usr/local/cuda/lib64'
4. unzip `TensorRT-7.2.1.6.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.0.tar.gz` at `/usr/local`, create softlink by `ln -s /usr/local/TensorRT-7.2.1.6 /usr/local/tensorRT`
5. add `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/tensorRT/lib:/usr/local/cuda/lib64:/usr/local/lib`, `export CPATH=$CPATH:/usr/local/cuda/include:/usr/local/tensorRT/include` to `~/.bashrc`
6. run `source ~/.bashrc`
>> CUDA 11.1 + TensorRT 7.2.1 for this repository (tested)
>>
>> CUDA 11.1 + TensorRT 8.5 for this repository (tested)
## How to generate trt model from onnx ?
```shell
trtexec --onnx=./vehicle.onnx --saveEngine=vehicleXXX.trt --buildOnly=true
```
## How to build trt_vehicle ?
we can build trt_vehicle separately.
0. set the right library path and include path for TensorRT in `CMakeLists.txt`
1. `mkdir build && cd build`
2. `cmake ..`
3. `make -j8`
all lib files saved to `build/libs`, all samples saved to `build/samples`. please refer to videopipe about how to run samples for trt_vehicle.
## Sample screenshot ##
### vehicle detect
![](../../doc/3rdparty/4.png)
### vehicle plate detect
![](../../doc/3rdparty/2.png)
### vehicle scan
![](../../doc/3rdparty/3.png)
### vehicle color and type classify
![](../../doc/3rdparty/1.png)
### vehicle compare
![](../../doc/3rdparty/5.png)
### vehicle search
![](../../doc/3rdparty/6.png)
### vehicle cluster
![](../../doc/3rdparty/7.png)

48
third_party/trt_vehicle/samples/vehicle.cpp vendored Normal file
View File

@@ -0,0 +1,48 @@
#include <opencv2/core/core.hpp>
#include <string>
#include <vector>
#include <map>
#include "../models/vehicle_plate_detector.h"
#include "../models/vehicle_detector.h"
using namespace std;
/*
* vehicle detect demo
*/
std::map<std::string, cv::Scalar> colors {{"car", cv::Scalar(255, 0, 0)},
{"bus", cv::Scalar(0, 255, 255)},
{"truck", cv::Scalar(0, 255, 0)}};
int main() {
std::string vehicle_model_path = "./vp_data/models/trt/vehicle/vehicle_v8.5.trt";
auto image = cv::imread("./vp_data/test_images/vehicle/0.jpg");
auto image2 = cv::imread("./vp_data/test_images/vehicle/1.jpg");
auto vehicle_detector = new trt_vehicle::VehicleDetector(vehicle_model_path);
std::vector<cv::Mat> img_datas = {image, image2};
std::vector<std::vector<trt_vehicle::ObjBox>> results;
vehicle_detector->detect(img_datas, results);
for (int i = 0; i < results.size(); i++)
{
/* code */
for (int j = 0; j < results[i].size(); j++)
{
/* code */
auto& objbox = results[i][j];
cv::rectangle(img_datas[i], cv::Rect(objbox.x, objbox.y, objbox.width, objbox.height), colors.at(objbox.label), 2);
cv::putText(img_datas[i], objbox.label + std::to_string(objbox.score), cv::Point(objbox.x, objbox.y), 1, 0.8, colors.at(objbox.label));
}
}
cv::imshow("image", image);
cv::imshow("image2", image2);
cv::waitKey(0);
delete vehicle_detector;
}

143
third_party/trt_vehicle/samples/vehicle_cluster.cpp vendored Normal file
View File

@@ -0,0 +1,143 @@
#include <opencv2/core/core.hpp>
#include <opencv2/freetype.hpp>
#include <string>
#include <vector>
#include <memory>
#include <experimental/filesystem>
#include "../models/vehicle_feature_encoder.h"
#include "../../../utils/vp_utils.h"
#include "../../bhtsne/tsne.h" // t-SNE algo
using namespace std;
/*
* vehicle cluster demo, reduce dims of features and display them on 2D screen.
*/
// extract feature for vehicle
void extract_feature(std::string vehicle_img_path, std::vector<float>& feature) {
static std::string feature_model_path = "./vp_data/models/trt/vehicle/vehicle_embedding_v8.5.trt";
static std::shared_ptr<trt_vehicle::VehicleFeatureEncoder> vehicleEncoder = nullptr;
if (!vehicleEncoder) {
vehicleEncoder = std::make_shared<trt_vehicle::VehicleFeatureEncoder>(feature_model_path);
}
auto vehicle = cv::imread(vehicle_img_path);
std::vector<std::vector<float>> results;
std::vector<cv::Mat> img_datas = {vehicle};
vehicleEncoder->encode(img_datas, results);
for(auto& i : results[0]) {
feature.push_back(i);
}
}
// load all vehicle images (.jpg) from disk and extract all features
void load_vehicle_dataset(std::string dataset_dir,
std::vector<std::pair<std::string, std::vector<float>>>& features_set) {
features_set.clear();
// iterate directory
using recursive_directory_iterator = std::experimental::filesystem::recursive_directory_iterator;
for (const auto& dir_entry : recursive_directory_iterator(dataset_dir))
if (vp_utils::ends_with(dir_entry.path(), ".jpg")) {
std::cout << "load vehicle image: " << dir_entry << std::endl;
// extract single feature
std::vector<float> feature;
extract_feature(dir_entry.path(), feature);
std::pair<std::string, std::vector<float>> p {dir_entry.path(), feature};
features_set.push_back(p);
}
}
void reduce_dims(std::vector<std::pair<std::string, std::vector<float>>>& features_set,
std::vector<std::pair<std::string, std::vector<float>>>& low_dims_features_set,
/* default parameters for t-SNE algorithm */
int no_dims = 2, int max_iter = 1000, double perplexity = 2, double theta = 0.5, int rand_seed = -1, bool skip_random_init = false, int stop_lying_iter = 250, int mom_switch_iter = 250) {
assert(features_set.size() != 0);
auto N = features_set.size();
auto D = features_set[0].second.size(); // all the same as the first feature's dims
// prepare input
double data[N * D];
double Y[N * no_dims];
for (int i = 0; i < N; i++) {
for (int j = 0; j < D; j++) {
data[i * D + j] = features_set[i].second[j];
}
}
// call t-SNE
TSNE::run(data, N, D, Y, no_dims, perplexity, theta, rand_seed, skip_random_init, max_iter, stop_lying_iter, mom_switch_iter);
// prepare output
for (int i = 0; i < N; i++) {
std::vector<float> low_dims_feature;
for (int j = 0; j < no_dims; j++) {
low_dims_feature.push_back(float(Y[i * no_dims + j]));
}
std::pair<std::string, std::vector<float>> p {features_set[i].first, low_dims_feature};
low_dims_features_set.push_back(p);
}
}
int main() {
auto vehicle_dataset_dir = "./vp_data/test_images/vehicle_feature";
auto ft2 = cv::freetype::createFreeType2();
ft2->loadFontData("./vp_data/font/NotoSansCJKsc-Medium.otf", 0);
// load vehicle dataset
std::vector<std::pair<std::string, std::vector<float>>> features_set;
load_vehicle_dataset(vehicle_dataset_dir, features_set);
// reduce dims using t-SNE
std::vector<std::pair<std::string, std::vector<float>>> low_dims_feature_set;
reduce_dims(features_set, low_dims_feature_set);
// print low dims features
std::cout << "low dims features:" << std::endl;
for(auto& i: low_dims_feature_set) {
std::cout << "[" << std::endl;
for(auto& j: i.second) {
std::cout << j << ",";
}
std::cout << "]" << std::endl;
}
// normalize low dims feature to coordinate of [0:1] and display them on 2D screen
auto max_x = 0.0f, max_y = 0.0f, min_x = 0.0f, min_y = 0.0f;
for(auto& i: low_dims_feature_set) {
auto& f = i.second;
// 2 values in f
max_x = std::max(max_x, f[0]);
max_y = std::max(max_y, f[1]);
min_x = std::min(min_x, f[0]);
min_y = std::min(min_y, f[1]);
}
auto x_range = max_x - min_x;
auto y_range = max_y - min_y;
// draw on (canvas_w_h + img_w_h) * (canvas_w_h + img_w_h)
auto canvas_w_h = 800, img_w_h = 70;
cv::Mat canvas(canvas_w_h + img_w_h, canvas_w_h + img_w_h, CV_8UC3, cv::Scalar(127, 127, 127));
for(auto& i: low_dims_feature_set) {
auto& f =i.second;
// convert to [0:1]
f[0] = (f[0] - min_x) / x_range;
f[1] = (f[1] - min_y) / y_range;
auto img = cv::imread(i.first);
cv::Mat img_tmp;
cv::resize(img, img_tmp, cv::Size(img_w_h, img_w_h));
cv::rectangle(img_tmp, cv::Rect(0, 0, img_tmp.cols, img_tmp.rows), cv::Scalar(255, 0, 0));
cv::Mat roi(canvas, cv::Rect(int(f[0] * canvas_w_h), int(f[1] * canvas_w_h), img_w_h, img_w_h));
img_tmp.copyTo(roi);
}
cv::imshow("cluster result", canvas);
cv::waitKey(0);
return 0;
}

88
third_party/trt_vehicle/samples/vehicle_color_type.cpp vendored Normal file
View File

@@ -0,0 +1,88 @@
#include <opencv2/core/core.hpp>
#include <opencv2/freetype.hpp>
#include <string>
#include <vector>
#include <map>
#include "../models/vehicle_color_classifier.h"
#include "../models/vehicle_type_classifier.h"
using namespace std;
/*
* vehicle color and vehicle type classifier demo
* tensorRT 8.5 + CUDA 11.1 + cuDNN 8.6 for this code, because original onnx weight could not be converted to TRT engine on tensorRT 7.2
*
* ## color ##
* black
* blue
* grey
* other
* red
* white
* yellow
*
* ## type ##
* bus
* business_car
* construction_truck
* large_truck
* sedan
* small_truck
* suv
* tanker
* van
*/
// to Chinese
std::map<std::string, std::string> types {{"bus", "巴士"},
{"business_car", "商务面包车"},
{"construction_truck", "施工车"},
{"large_truck", "大卡车"},
{"sedan", "轿车"},
{"small_truck", "小卡车"},
{"suv", "SUV"},
{"tanker", "罐车"},
{"van", "厢式货车"}};
int main() {
std::string vehicle_color_model_path = "./vp_data/models/trt/vehicle/vehicle_color_v8.5.trt";
std::string vehicle_type_model_path = "./vp_data/models/trt/vehicle/vehicle_type_v8.5.trt";
std::string font_path = "./vp_data/font/NotoSansCJKsc-Medium.otf";
auto image = cv::imread("./vp_data/test_images/vehicle_cls/1.jpg"); // vehicle in the center of image, crop it first if possible
auto image2 = cv::imread("./vp_data/test_images/vehicle_cls/2.jpg");
auto image3 = cv::imread("./vp_data/test_images/vehicle_cls/3.jpg");
auto image4 = cv::imread("./vp_data/test_images/vehicle_cls/4.jpg");
auto image5 = cv::imread("./vp_data/test_images/vehicle_cls/5.jpg");
auto image6 = cv::imread("./vp_data/test_images/vehicle_cls/6.jpg");
cv::Ptr<cv::freetype::FreeType2> ft2 = cv::freetype::createFreeType2();
ft2->loadFontData(font_path, 0);
auto vehicle_color_classifier = new trt_vehicle::VehicleColorClassifier(vehicle_color_model_path);
auto vehicle_type_classifier = new trt_vehicle::VehicleTypeClassifier(vehicle_type_model_path);
std::vector<cv::Mat> img_datas = {image, image2, image3, image4, image5, image6};
std::vector<trt_vehicle::ObjCls> color_results;
std::vector<trt_vehicle::ObjCls> type_results;
vehicle_color_classifier->classify(img_datas, color_results); // color classify
vehicle_type_classifier->classify(img_datas, type_results); // type classify
for (int i = 0; i < img_datas.size(); i++) {
/* draw class label at top left */
// cv::putText(img_datas[i], color_results[i].label + " " + types.at(type_results[i].label), cv::Point(30, 30), 1, 2, colors.at(color_results[i].label));
ft2->putText(img_datas[i], color_results[i].label + " " + types.at(type_results[i].label), cv::Point(30, 30), 25, cv::Scalar(0, 255, 0), cv::FILLED, cv::LINE_AA, true);
}
cv::imshow("image", image);
cv::imshow("image2", image2);
cv::imshow("image3", image3);
cv::imshow("image4", image4);
cv::imshow("image5", image5);
cv::imshow("image6", image6);
cv::waitKey(0);
delete vehicle_color_classifier;
delete vehicle_type_classifier;
}

97
third_party/trt_vehicle/samples/vehicle_compare.cpp vendored Normal file
View File

@@ -0,0 +1,97 @@
#include <opencv2/core/core.hpp>
#include <opencv2/freetype.hpp>
#include <string>
#include <vector>
#include "../models/vehicle_feature_encoder.h"
using namespace std;
/*
* vehicle comparision demo, 1:1
*/
// compare two features
// dis_type 0 means cosine distance, bigger means more similiar
// dis_type 1 means L2 distance, smaller means more similiar
double match(std::vector<float>& feature1, std::vector<float>& feature2, int dis_type) {
auto _face_feature1 = cv::Mat(1, feature1.size(), CV_32F, feature1.data());
auto _face_feature2 = cv::Mat(1, feature2.size(), CV_32F, feature2.data());
cv::normalize(_face_feature1, _face_feature1);
cv::normalize(_face_feature2, _face_feature2);
if(dis_type == 0) {
return cv::sum(_face_feature1.mul(_face_feature2))[0];
}
else if(dis_type == 1) {
return cv::norm(_face_feature1, _face_feature2);
}
else {
return 0;
}
}
int main() {
std::string featureModelPath = "./vp_data/models/trt/vehicle/vehicle_embedding_v8.5.trt";
auto image = cv::imread("./vp_data/test_images/vehicle_feature/2_001.jpg"); // first vehicle
auto image2 = cv::imread("./vp_data/test_images/vehicle_feature/2_003.jpg"); // second vehicle to compare(the same vehicle with different angle)
auto image3 = cv::imread("./vp_data/test_images/vehicle_feature/5_002.jpg"); // third vehicle to compare(not the same vehicle)
auto vehicleEncoder = new trt_vehicle::VehicleFeatureEncoder(featureModelPath);
auto ft2 = cv::freetype::createFreeType2();
ft2->loadFontData("./vp_data/font/NotoSansCJKsc-Medium.otf", 0);
std::vector<std::vector<float>> results;
std::vector<cv::Mat> img_datas = {image, image2, image3};
vehicleEncoder->encode(img_datas, results); // encode
auto feature1 = results[0];
auto feature2 = results[1];
auto feature3 = results[2];
// print features, 256 dims
auto print_features = [](std::vector<float> features, std::string name) {
std::cout << name << ": " << features.size() << " dims " << "[" << std::endl;
for (auto& i: features) {
std::cout << i << ",";
}
std::cout << "]" << std::endl;
};
print_features(feature1, "feature1");
print_features(feature2, "feature2");
print_features(feature3, "feature3");
auto dis_type = 0;
auto dis_1_2 = match(feature1, feature2, dis_type); // compare feature1 and feature2
auto dis_1_3 = match(feature1, feature3, dis_type); // compare feature1 and feature3
std::cout << "distance between feature1 and feature2: " << dis_1_2 << std::endl;
std::cout << "distance between feature1 and feature3: " << dis_1_3 << std::endl;
// convert to similarity
auto similiarity_1_2 = dis_type == 0 ? dis_1_2 : (1 - dis_1_2);
similiarity_1_2 = std::max(similiarity_1_2, 0.0);
auto similiarity_1_3 = dis_type == 0 ? dis_1_3 : (1 - dis_1_3);
similiarity_1_3 = std::max(similiarity_1_3, 0.0);
// display
auto h = image.rows + std::max(image2.rows, image3.rows) + 50 + 10 + 10;
auto w = std::max(image2.cols, image3.cols) * 2 + 50 + 10 + 10;
cv::Mat canvas(h, w, CV_8UC3,cv::Scalar(127, 127, 127));
cv::Mat roi_(canvas,cv::Rect(canvas.cols / 2 - image.cols / 2, 10, image.cols, image.rows));
cv::Mat roi_2(canvas, cv::Rect(canvas.cols / 2 - image2.cols - 25, 10 + image.rows + 25, image2.cols, image2.rows));
cv::Mat roi_3(canvas, cv::Rect(canvas.cols / 2 + 25, 10 + image.rows + 25, image3.cols, image3.rows));
image.copyTo(roi_);
image2.copyTo(roi_2);
image3.copyTo(roi_3);
cv::line(canvas, cv::Point(canvas.cols / 2, 10 + image.rows), cv::Point(10 + image2.cols / 2, 10 + image.rows + 50), cv::Scalar(255, 0, 0));
cv::line(canvas, cv::Point(canvas.cols / 2, 10 + image.rows), cv::Point(canvas.cols / 2 + 25 + image3.cols / 2, 10 + image.rows + 50), cv::Scalar(255, 0, 0));
ft2->putText(canvas, "相似度：" + std::to_string(similiarity_1_2), cv::Point(10, 10 + image.rows + 25), 20, cv::Scalar(255, 0, 0), cv::FILLED, cv::LINE_AA, true);
ft2->putText(canvas, "相似度：" + std::to_string(similiarity_1_3), cv::Point(10 + image2.cols + 50, 10 + image.rows + 25), 20, cv::Scalar(255, 0, 0), cv::FILLED, cv::LINE_AA, true);
cv::imshow("compare result", canvas);
cv::waitKey(0);
delete vehicleEncoder;
return 0;
}

52
third_party/trt_vehicle/samples/vehicle_plate.cpp vendored Executable file
View File

@@ -0,0 +1,52 @@
#include <opencv2/core/core.hpp>
#include <opencv2/freetype.hpp>
#include <string>
#include <vector>
#include "../models/vehicle_plate_detector.h"
using namespace std;
/*
* vehicle plate demo
*/
std::map<std::string, cv::Scalar> colors {{"blue", cv::Scalar(255, 0, 0)},
{"green", cv::Scalar(0, 255, 0)},
{"yellow", cv::Scalar(0, 255, 255)},
{"white", cv::Scalar(255, 255, 255)}};
int main(){
std::string plateModelPath = "./vp_data/models/trt/plate/det_v8.5.trt";
std::string charModelPath = "./vp_data/models/trt/plate/rec_v8.5.trt";
auto image = cv::imread("./vp_data/test_images/plates/0.jpg");
auto image2 = cv::imread("./vp_data/test_images/plates/1.jpg");
auto image3 = cv::imread("./vp_data/test_images/plates/2.jpg");
auto plateDetector = new trt_vehicle::VehiclePlateDetector(plateModelPath, charModelPath);
auto ft2 = cv::freetype::createFreeType2();
ft2->loadFontData("./vp_data/font/NotoSansCJKsc-Medium.otf", 0);
std::vector<std::vector<trt_vehicle::Plate>> results;
std::vector<cv::Mat> img_datas = {image, image2, image3};
plateDetector->detect(img_datas, results, true);
for (int i = 0; i < results.size(); i++) {
/* code */
for (int j = 0; j < results[i].size(); j++)
{
/* code */
auto& plate = results[i][j];
cv::rectangle(img_datas[i], cv::Rect(plate.x, plate.y, plate.width, plate.height), colors.at(plate.color), 2);
//cv::putText(img_datas[i], plate.text, cv::Point(plate.x, plate.y), 1, 0.8, colors.at(plate.color));
ft2->putText(img_datas[i], plate.text, cv::Point(plate.x, plate.y), 20, colors.at(plate.color), cv::FILLED, cv::LINE_AA, true);
}
}
cv::imshow("image", image);
cv::imshow("image2", image2);
cv::imshow("image3", image3);
cv::waitKey(0);
delete plateDetector;
return 0;
}

53
third_party/trt_vehicle/samples/vehicle_scan.cpp vendored Normal file
View File

@@ -0,0 +1,53 @@
#include <opencv2/core/core.hpp>
#include <string>
#include <vector>
#include <map>
#include "../models/vehicle_scanner.h"
using namespace std;
/*
* vehicle scan demo, detect on vehicle body
* tensorRT 8.5 + CUDA 11.1 + cuDNN 8.6 for this code, because original onnx weight could not be converted to TRT engine on tensorRT 7.2
*/
std::map<std::string, cv::Scalar> colors {{"truck", cv::Scalar(255, 0, 0)}, // 货车
{"bus", cv::Scalar(0, 255, 255)}, //巴士
{"car", cv::Scalar(0, 255, 0)}, //轿车
{"suv", cv::Scalar(0, 255, 0)}, //suv+商务
{"van", cv::Scalar(0, 255, 0)}, //面包+金杯
{"nowindow", cv::Scalar(0, 255, 0)}, //封窗货车，货拉拉
{"iveco", cv::Scalar(0, 255, 0)}, //9座依维柯
{"pickup", cv::Scalar(0, 255, 0)}, //皮卡
{"wheel", cv::Scalar(0, 255, 0)}}; //轮子
int main() {
std::string vehicle_scan_model_path = "./vp_data/models/trt/vehicle/vehicle_scan_v8.5.trt";
auto image = cv::imread("./vp_data/test_images/body/0.jpg");
auto image2 = cv::imread("./vp_data/test_images/body/1.jpg");
auto vehicle_scanner = new trt_vehicle::VehicleScanner(vehicle_scan_model_path);
std::vector<cv::Mat> img_datas = {image, image2};
std::vector<std::vector<trt_vehicle::ObjBox>> results;
vehicle_scanner->detect(img_datas, results);
for (int i = 0; i < results.size(); i++)
{
/* code */
for (int j = 0; j < results[i].size(); j++)
{
/* code */
auto& objbox = results[i][j];
cv::rectangle(img_datas[i], cv::Rect(objbox.x, objbox.y, objbox.width, objbox.height), colors.at(objbox.label), 2);
cv::putText(img_datas[i], objbox.label + std::to_string(objbox.score), cv::Point(objbox.x, objbox.y), 1, 0.8, colors.at(objbox.label));
}
}
cv::imshow("image", image);
cv::imshow("image2", image2);
cv::waitKey(0);
delete vehicle_scanner;
}

150
third_party/trt_vehicle/samples/vehicle_search.cpp vendored Normal file
View File

@@ -0,0 +1,150 @@
#include <opencv2/core/core.hpp>
#include <opencv2/freetype.hpp>
#include <string>
#include <vector>
#include <memory>
#include <experimental/filesystem>
#include "../models/vehicle_feature_encoder.h"
#include "../../../utils/vp_utils.h"
using namespace std;
/*
* vehicle search demo, 1:N
*/
// compare two features
// dis_type 0 means cosine distance, bigger means more similiar
// dis_type 1 means L2 distance, smaller means more similiar
double match(std::vector<float>& feature1, std::vector<float>& feature2, int dis_type) {
auto _face_feature1 = cv::Mat(1, feature1.size(), CV_32F, feature1.data());
auto _face_feature2 = cv::Mat(1, feature2.size(), CV_32F, feature2.data());
cv::normalize(_face_feature1, _face_feature1);
cv::normalize(_face_feature2, _face_feature2);
if(dis_type == 0) {
return cv::sum(_face_feature1.mul(_face_feature2))[0];
}
else if(dis_type == 1) {
return cv::norm(_face_feature1, _face_feature2);
}
else {
return 0;
}
}
// extract feature for vehicle
void extract_feature(std::string vehicle_img_path, std::vector<float>& feature) {
static std::string feature_model_path = "./vp_data/models/trt/vehicle/vehicle_embedding_v8.5.trt";
static std::shared_ptr<trt_vehicle::VehicleFeatureEncoder> vehicleEncoder = nullptr;
if (!vehicleEncoder) {
vehicleEncoder = std::make_shared<trt_vehicle::VehicleFeatureEncoder>(feature_model_path);
}
auto vehicle = cv::imread(vehicle_img_path);
std::vector<std::vector<float>> results;
std::vector<cv::Mat> img_datas = {vehicle};
vehicleEncoder->encode(img_datas, results);
for(auto& i : results[0]) {
feature.push_back(i);
}
}
// load all vehicle images (.jpg) from disk and extract all features
void load_vehicle_dataset(std::string dataset_dir,
std::vector<std::pair<std::string, std::vector<float>>>& features_set) {
features_set.clear();
// iterate directory
using recursive_directory_iterator = std::experimental::filesystem::recursive_directory_iterator;
for (const auto& dir_entry : recursive_directory_iterator(dataset_dir))
if (vp_utils::ends_with(dir_entry.path(), ".jpg")) {
std::cout << "load vehicle image: " << dir_entry << std::endl;
// extract single feature
std::vector<float> feature;
extract_feature(dir_entry.path(), feature);
std::pair<std::string, std::vector<float>> p {dir_entry.path(), feature};
features_set.push_back(p);
}
}
// match features using query feature, sorted by similiarity
void search(std::vector<float>& query_feature,
std::vector<std::pair<std::string, std::vector<float>>>& features_set,
std::vector<std::pair<std::string, float>>& query_result,
int dis_type = 0,
int top_n = 0) {
query_result.clear();
// just loop the features set
for (auto& i: features_set) {
/* code */
auto dis = match(i.second, query_feature, dis_type);
std::pair<std::string, float> p {i.first, dis};
query_result.push_back(p);
}
// sort from high to low
sort(query_result.begin(), query_result.end(), [=](std::pair<std::string, float>& a, std::pair<std::string, float>& b)
{
return a.second > b.second;
});
}
int main() {
auto vehicle_dataset_dir = "./vp_data/test_images/vehicle_feature";
auto query_vehicle_path = "./vp_data/test_images/vehicle_feature/7_002.jpg";
auto ft2 = cv::freetype::createFreeType2();
ft2->loadFontData("./vp_data/font/NotoSansCJKsc-Medium.otf", 0);
// load vehicle dataset
std::vector<std::pair<std::string, std::vector<float>>> features_set;
load_vehicle_dataset(vehicle_dataset_dir, features_set);
// load query vehicle
std::vector<float> query_feature;
extract_feature(query_vehicle_path, query_feature);
// search it!
std::vector<std::pair<std::string, float>> query_result;
search(query_feature, features_set, query_result, 0, 10);
// print similiarity from high to low
for(auto& i: query_result) {
std::cout << i.second << " ==> " << i.first << std::endl;
}
// display according to query_result
auto n_query = query_result.size();
auto rect_w_h = 80; auto gap = 20; auto cols = 10;
auto rows = n_query / cols + 2;
// create canvas
cv::Mat canvas(rows * (rect_w_h + gap) + gap, cols * (rect_w_h + gap) + gap, CV_8UC3,cv::Scalar(127, 127, 127));
// query vehicle at first line
cv::Mat roi_query = cv::Mat(canvas, cv::Rect(gap, gap, rect_w_h, rect_w_h));
auto query_vehicle_img = cv::imread(query_vehicle_path);
cv::Mat query_vehicle_img_tmp;
cv::resize(query_vehicle_img, query_vehicle_img_tmp, cv::Size(rect_w_h, rect_w_h));
query_vehicle_img_tmp.copyTo(roi_query);
ft2->putText(canvas, "query vehicle:", cv::Point(20, 14), 14, cv::Scalar(255, 0, 0), cv::FILLED, cv::LINE_AA, true);
// query result
for(int i = 0; i < query_result.size(); ++i) {
auto row = i / cols + 1;
auto col = i % cols;
cv::Mat roi = cv::Mat(canvas, cv::Rect(gap + col * (rect_w_h + gap), gap + row * (rect_w_h + gap), rect_w_h, rect_w_h));
auto vehicle_img = cv::imread(query_result[i].first);
cv::Mat vehicle_img_tmp;
cv::resize(vehicle_img, vehicle_img_tmp, cv::Size(rect_w_h, rect_w_h));
vehicle_img_tmp.copyTo(roi);
ft2->putText(canvas, std::to_string(std::max(query_result[i].second, 0.0f)), cv::Point(gap + col * (rect_w_h + gap), row * (rect_w_h + gap) + gap), 14, cv::Scalar(255, 0, 0), cv::FILLED, cv::LINE_AA, true);
}
cv::imshow("search result", canvas);
cv::waitKey(0);
return 0;
}

531
third_party/trt_vehicle/util/algorithm_util.cpp vendored Executable file
View File

@@ -0,0 +1,531 @@
#include "algorithm_util.h"
namespace trt_vehicle{
float overlap( float x1, float w1, float x2, float w2)
{
float l1 = x1;
float l2 = x2;
float left = ( l1 > l2 ) ? l1 : l2;
float r1 = x1 + w1;
float r2 = x2 + w2;
float right = ( r1 < r2 ) ? r1 : r2;
return ( right - left );
}
//calculation intersection area
float boxIntersection(ObjBox& a, ObjBox& b)
{
float w = overlap( a.x, a.width, b.x, b.width);
float h = overlap( a.y, a.height, b.y, b.height);
if( w < 0 || h < 0 )
{
return 0.0;
}
return ( w * h );
}
//calculation union area
float boxUnion(ObjBox& a, ObjBox& b)
{
float comm_area = boxIntersection( a, b );
return ( a.width * a.height + b.width * b.height - comm_area );
}
//calculate iou of tow boxes
float rectIou(ObjBox& a, ObjBox& b)
{
return boxIntersection(a, b)/boxUnion(a, b);
}
//calculate iou of tow boxes
float rectIouRegion(ObjBox& a, ObjBox& b)
{
int inner = 0;
for(int i=0;i<a.gridCoords.size();i++){
for(int j=0;j<b.gridCoords.size();j++){
if(a.gridCoords[i].i == b.gridCoords[j].i){
inner++;
break;
}
}
}
return inner*1.0/(a.gridCoords.size()+b.gridCoords.size()-inner);
}
//calculate iou of tow boxes
float rectIouDivOne(ObjBox& a, ObjBox& b)
{
float inner = boxIntersection(a, b)*1.0;
return max(inner/(a.width * a.height),inner/(b.width * b.height));
}
//coords change : from xyxy to objBox
void xyxy2objBox(int x1,int y1,int x2,int y2, ObjBox& b )
{
b.x = x1;
b.width = x2-x1;
b.y = y1;
b.height = y2-y1;
}
//Sort from large to small , return indeices
std::vector<int> getSortIndex(std::vector<float>& scores)
{
std::vector<int> indices;
for (int i = 0; i < scores.size(); i++)
indices.push_back(i);
for (int i = 0; i < scores.size(); i++)
for (int j = i + 1; j < scores.size(); j++)
{
if(scores[j] > scores[i])
{
int tmp = indices[i];
indices[i] = indices[j];
indices[j] = tmp;
float tmpS = scores[i];
scores[i] = scores[j];
scores[j] = tmpS;
}
}
return indices;
}
//nms for bounding boxes
int nonMaximumSuppression(vector<ObjBox>& boxes, vector<float>& scores, float overlapThreshold,vector<ObjBox>& outBoxes)
{
outBoxes.clear();
std::vector<int> indices = getSortIndex(scores);
int numBoxes = scores.size();
vector<float> box_area(numBoxes);
vector<bool> is_suppressed(numBoxes);
for (int i = 0; i < numBoxes; i++)
{
is_suppressed[i] = false;
box_area[i] = (float)(boxes[i].width*boxes[i].height);
}
for (int i = 0; i < numBoxes; i++)
{
if (!is_suppressed[indices[i]])
{
outBoxes.push_back(boxes[indices[i]]);
for (int j = i + 1; j < numBoxes; j++)
{
if (!is_suppressed[indices[j]])
{
if (rectIou(boxes[indices[i]],boxes[indices[j]]) > overlapThreshold)
{
is_suppressed[indices[j]] = true;
}
}
}
}
}
return true;
}
//difference map pf two pictures
cv::Mat diffImgs(cv::Mat bgImg,cv::Mat frontImg){
cv::Mat bgGray,frontGray,binary,diff;
if (bgImg.channels() == 3)
cv::cvtColor(bgImg, bgGray, cv::COLOR_BGR2GRAY);
else
bgGray = bgImg.clone();
if (frontImg.channels() == 3)
cv::cvtColor(frontImg, frontGray, cv::COLOR_BGR2GRAY);
else
frontGray = frontImg.clone();
cv::absdiff(bgGray, frontGray, diff);
cv::threshold(diff, binary, 5, 255, 0);
cv::Mat kernel5 = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));
cv::Mat kernel7 = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(7, 7));
cv::morphologyEx(binary, binary, cv::MORPH_OPEN, kernel5);
cv::morphologyEx(binary, binary, cv::MORPH_CLOSE, kernel7);
return binary;
}
//sort
void sort(std::vector<float>& scores,std::vector<float>& coordXs,std::vector<float>& coordYs)
{
for (int i = 0; i < scores.size(); i++)
for (int j = i + 1; j < scores.size(); j++)
{
if(scores[j] < scores[i])
{
float s = scores[i];
scores[i] = scores[j];
scores[j] = s;
float x = coordXs[i];
coordXs[i] = coordXs[j];
coordXs[j] = x;
float y = coordYs[i];
coordYs[i] = coordYs[j];
coordYs[j] = y;
}
}
}
//sort
void sort(std::vector<double>& scores)
{
for (int i = 0; i < scores.size(); i++)
for (int j = i + 1; j < scores.size(); j++)
{
if(scores[j] < scores[i])
{
double s = scores[i];
scores[i] = scores[j];
scores[j] = s;
}
}
}
//draw boxes to img , and return the drawed img
cv::Mat drawRectangle(cv::Mat img,std::vector<ObjBox>& boxes,cv::Scalar color){
cv::Mat imgDraw = img.clone();
for(int i=0;i<boxes.size();i++){
ObjBox box = boxes[i];
cv::Rect rect=cv::Rect(int(box.x),int(box.y),int(box.width),int(box.height));
cv::rectangle(imgDraw,rect,color,1);
}
return imgDraw;
}
//Calculate the Euclidean distance between two vectors
double calVectorDistance(std::vector<float> f1,std::vector<float> f2){
double dDisTmp = 0.0;
double eDisTmp1 = 0.0;
double eDisTmp2 = 0.0;
for(int i=0;i<f1.size();i++){
eDisTmp1 += f1[i]*f1[i];
eDisTmp2 += f2[i]*f2[i];
}
double eDis1 = sqrt(eDisTmp1);
double eDis2 = sqrt(eDisTmp2);
for(int i=0;i<f1.size();i++){
dDisTmp += pow((f1[i]/eDis1-f2[i]/eDis2),2);
}
double dis = sqrt(dDisTmp);
return dis;
}
//Batch calculation of Euclidean distance between two vectors
std::vector<double> calVectorDistances(std::vector<std::vector<float>> fs1,std::vector<std::vector<float>> fs2){
std::vector<double> distants;
for(int i=0;i<fs1.size();i++){
distants.push_back(calVectorDistance(fs1[i],fs2[i]));
}
return distants;
}
//Generate binary vector from distance value(score)
void genBinaryVectorfromScore(vector<int> binaryScores,vector<double> scores,int imgW,int imgH,int downScale){
binaryScores.clear();
int gridW = int(imgW/downScale);
int gridH = int(imgH/downScale);
for(int i=0;i<gridH*gridW;i++){
int score = std::min(255,int(scores[i]*255));
if(score > 155){
binaryScores.push_back(1);
}else {
binaryScores.push_back(0);
}
}
}
//Generate gray image from distance value(score)
void genImgfromScore(cv::Mat& img,vector<double> scores,int imgW,int imgH,int downScale){
img = img*0;
int gridW = int(imgW/downScale);
int gridH = int(imgH/downScale);
for(int i=0;i<gridH;i++){
for(int j=0;j<gridW;j++){
int score = std::min(255,int(scores[i*gridW+j]*255));
if(score > 100){
cv::Rect roi_rect = cv::Rect(j*downScale, i*downScale, downScale,downScale);
img(roi_rect).setTo(score);
}
}
}
}
//filter rectangles through regions of interest region
std::vector<ObjBox> selectRectByInterestMask(std::vector<ObjBox>& rects,cv::Mat interestedMask){
std::vector<ObjBox> interestedRects;
interestedRects.clear();
for(int i=0;i<rects.size();i++){
ObjBox box = rects[i];
cv::Rect rect=cv::Rect(int(box.x),int(box.y),int(box.width),int(box.height));
cv::Mat imgSeg = interestedMask(rect);
int noZeroCount = cv::countNonZero(imgSeg);
float ratio = noZeroCount*1.0/(box.width*box.height);
if(ratio>0.5){
interestedRects.push_back(box);
}
}
return interestedRects;
}
//find adjacent grids
int disInVector(GridCoord gridCoord,std::vector<std::vector<GridCoord>>& indicesAll){
int index = -1;
for(int i=0;i<indicesAll.size();i++){
for(int j=0;j<indicesAll[i].size();j++){
if(abs(gridCoord.x-indicesAll[i][j].x)<=1 && abs(gridCoord.y-indicesAll[i][j].y)<=1){
return i;
}
}
}
return index;
}
//generate rectangles using distance vector
void genBoxesFromVector(std::vector<double>& dis,int gridW,int downScale,std::vector<ObjBox>* rects,std::vector<std::vector<GridCoord>>* indicesAll){
for(int i=0; i<dis.size(); i++){
if(dis[i] == 0){
continue;
}
int x = i%gridW;
int y = i/gridW;
GridCoord gridCoord;
gridCoord.x = x;
gridCoord.y = y;
gridCoord.i = i;
int index = disInVector(gridCoord,*indicesAll);
if(index >= 0){
(*indicesAll)[index].push_back(gridCoord);
}else{
std::vector<GridCoord> indicesSingle;
indicesSingle.push_back(gridCoord);
(*indicesAll).push_back(indicesSingle);
}
}
for(int i=0;i<indicesAll->size();i++){
ObjBox rect;
int minX=9999;
int minY=9999;
int maxX=-1;
int maxY=-1;
for(int j=0;j<(*indicesAll)[i].size();j++){
if(minX>(*indicesAll)[i][j].x){
minX = (*indicesAll)[i][j].x;
}
if(minY>(*indicesAll)[i][j].y){
minY = (*indicesAll)[i][j].y;
}
if(maxX<(*indicesAll)[i][j].x){
maxX = (*indicesAll)[i][j].x;
}
if(maxY<(*indicesAll)[i][j].y){
maxY = (*indicesAll)[i][j].y;
}
}
rect.x = minX*downScale;
rect.y = minY*downScale;
rect.width = (maxX-minX+1)*downScale;
rect.height = (maxY-minY+1)*downScale;
rects->push_back(rect);
}
}
//generate rectangles using binary
std::vector<ObjBox> genBoxesFromBinary(cv::Mat& binary){
std::vector<vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(binary, contours, hierarcy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE, cv::Point());
std::vector<ObjBox> rects;
rects.clear();
for( int i = 0 ; i < contours.size(); i++){
cv::Rect rect = cv::boundingRect(cv::Mat(contours[i]));
ObjBox r;
r.x = rect.x;
r.y = rect.y;
r.width = rect.width;
r.height = rect.height;
rects.push_back(r);
}
return rects;
}
//compare the current image with the background image to find different areas
bool compareAlarmRectToBg(ObjBox rectCurr,std::vector<ObjBox>& rects,ObjBox& rect){
float iouMax = 0.0;
int index = -1;
for( int i = 0 ; i < rects.size(); i++){
ObjBox rect = rects[i];
float iou = rectIou(rectCurr,rect);
if(iou > iouMax){
iouMax = iou;
index = i;
}
}
if(iouMax > 0.25){
rect = rects[index];
rects.erase( rects.begin() + index);
return true;
}
return false;
}
//filter rectangular boxes by scale / aspect ratio / region of interest / iou , etc
void RectSelector::init(std::vector<ObjBox> rects){
m_rects.clear();
m_rects = rects;
m_cursor = -1;
}
ObjBox RectSelector::next(){
m_cursor = m_cursor+1;
return m_rects[m_cursor];
}
bool RectSelector::hasNext(){
return m_rects.size()> m_cursor+1;
}
void RectSelector::remove(){
m_rects.erase( m_rects.begin() + m_cursor );
m_cursor = m_cursor-1;
}
void RectSelector::remove_inside(ObjBox rect){
for( int i = 0 ; i < m_rects.size(); i++){
if(i == m_cursor){continue;}
ObjBox r = m_rects[i];
if(rect.y>=r.y && rect.y+rect.height<=r.y+r.height && rect.x>=r.x && rect.x+rect.width<=r.x+r.width){
remove();
break;
}
}
}
void RectSelector::remove_overlop(ObjBox rect,float iouThres){
for( int i = 0 ; i < m_rects.size(); i++){
if(i == m_cursor){continue;}
ObjBox r = m_rects[i];
float boxIou = rectIou(r, rect);
if(boxIou>iouThres){
remove();
break;
}
}
}
void RectSelector::rmInside(){
m_cursor = -1;
while(hasNext()){
ObjBox rect = next();
remove_inside(rect);
}
}
void RectSelector::rmOverlop(float iouThres){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
remove_overlop(rect,iouThres);
}
}
void RectSelector::rmLittle(int minW,int minH){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
if (rect.width <= minW || rect.height <= minH)
{ remove();}
}
}
void RectSelector::rmBig(int maxW,int maxH){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
if (rect.width >= maxW || rect.height >= maxH)
{
remove();
}
}
}
void RectSelector::rmBigRatio(int maxW,int maxH,int imgH, float ratio){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
int centerY = rect.y+rect.height/2;
float currRatio = ratio + centerY*1.0/imgH*(1-ratio);
if (rect.width >= maxW*currRatio || rect.height >= maxH*currRatio)
{
remove();
}
}
}
void RectSelector::rmBigWithRoadMask(cv::Mat& mask,std::vector<float>& roadNumsList,float carLeastWidth){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
int centerX,centerY;
centerX = rect.x+rect.width/2;
centerY = rect.y+rect.height/2;
cv::Rect rectRoi=cv::Rect(0, int(centerY-1), mask.size().width, 2);
cv::Mat roiMask = mask(rectRoi);
std::vector<vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(roiMask, contours, hierarcy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE, cv::Point());
int index = -1;
int roadWidth = -1;
std::vector<cv::Rect> rectMulti;
for( int i = 0 ; i < contours.size(); i++){
cv::Rect rectSingle = cv::boundingRect(cv::Mat(contours[i]));
rectMulti.push_back(rectSingle);
}
for( int i = 0 ; i < rectMulti.size(); i++){
for( int j = i+1 ; j < rectMulti.size(); j++){
if(rectMulti[i].x>rectMulti[j].x){
cv::Rect tmp = rectMulti[i];
rectMulti[i] = rectMulti[j];
rectMulti[j] = tmp;
}
}
}
for( int i = 0 ; i < rectMulti.size(); i++){
cv::Rect rectSingle = rectMulti[i];
if(rectSingle.x<=centerX && centerX<=rectSingle.x+rectSingle.width){
index = i;
roadWidth = rectSingle.width;
}
}
if(index>=roadNumsList.size() || index == -1){
continue;
}
float carWidth = max(roadWidth/roadNumsList[index],carLeastWidth);
if (rect.width > carWidth || rect.height > carWidth){
remove();
}
}
}
void RectSelector::rmRatio(){
m_cursor = -1 ;
while(hasNext()){
ObjBox rect = next();
float ratio = rect.width*1.0/rect.height;
float thres = 4.0;
if (ratio >= thres || ratio <= 1.0/thres)
{ remove();}
}
}
std::vector<ObjBox> RectSelector::getRect(){
return m_rects;
}
}

120
third_party/trt_vehicle/util/algorithm_util.h vendored Executable file
View File

@@ -0,0 +1,120 @@
#ifndef __IMAGE_ANALYSIS_ROI_OBJECT_BOX_UTIL_H__
#define __IMAGE_ANALYSIS_ROI_OBJECT_BOX_UTIL_H__
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <string>
#include <vector>
using namespace std;
namespace trt_vehicle{
struct GridCoord
{
int x;
int y;
int i;
};
struct ObjBox
{
int x;
int y;
int width;
int height;
float score;
int class_;
std::string label;
std::vector<GridCoord> gridCoords;
};
struct ObjCls {
int class_;
std::string label;
float score;
};
float overlap( float x1, float w1, float x2, float w2);
//calculation intersection area
float boxIntersection(ObjBox& a, ObjBox& b);
//calculation union area
float boxUnion(ObjBox& a, ObjBox& b);
//calculate iou of tow boxes
float rectIou(ObjBox& a, ObjBox& b);
float rectIouRegion(ObjBox& a, ObjBox& b);
float rectIouDivOne(ObjBox& a, ObjBox& b);
//coords change : from xyxy to objBox
void xyxy2objBox(int x1,int y1,int x2,int y2, ObjBox& b);
//Sort from large to small , return indeices
std::vector<int> getSortIndex(std::vector<float>& scores);
//nms for bounding boxes
int nonMaximumSuppression(vector<ObjBox>& boxes, vector<float>& scores, float overlapThreshold,vector<ObjBox>& outBoxes);
//difference map pf two pictures
cv::Mat diffImgs(cv::Mat bgImg,cv::Mat frontImg);
//sort
void sort(std::vector<float>& scores,std::vector<float>& coordXs,std::vector<float>& coordYs);
void sort(std::vector<double>& scores);
//draw boxes to img , and return the drawed img
cv::Mat drawRectangle(cv::Mat img,std::vector<ObjBox>& boxes,cv::Scalar color);
//Calculate the Euclidean distance between two vectors
double calVectorDistance(std::vector<float> f1,std::vector<float> f2);
//Batch calculation of Euclidean distance between two vectors
std::vector<double> calVectorDistances(std::vector<std::vector<float>> fs1,std::vector<std::vector<float>> fs2);
//Generate binary vector from distance value(score)
void genBinaryVectorfromScore(vector<int> binaryScores,vector<double> scores,int imgW,int imgH,int downScale);
//Generate gray image from distance value(score)
void genImgfromScore(cv::Mat& img,vector<double> scores,int imgW,int imgH,int downScale);
//filter rectangles through regions of interest region
std::vector<ObjBox> selectRectByInterestMask(std::vector<ObjBox>& rects,cv::Mat interestedMask);
//find adjacent grids
int disInVector(GridCoord gridCoord,std::vector<std::vector<GridCoord>>& indicesAll);
//generate rectangles using distance vector
void genBoxesFromVector(std::vector<double>& dis,int gridW,int downScale,std::vector<ObjBox>* rects,std::vector<std::vector<GridCoord>>* indicesAll);
//generate rectangles using binary
std::vector<ObjBox> genBoxesFromBinary(cv::Mat& binary);
//compare the current image with the background image to find different areas
bool compareAlarmRectToBg(ObjBox rectCurr,std::vector<ObjBox>& rects,ObjBox& rect);
//filter rectangular boxes by scale / aspect ratio / region of interest / iou , etc
class RectSelector{
private:
std::vector<ObjBox> m_rects;
int m_cursor;
ObjBox next();
bool hasNext();
void remove();
void remove_inside(ObjBox rect);
void remove_overlop(ObjBox rect,float iouThres);
public:
void init(std::vector<ObjBox> rects);
void rmInside();
void rmOverlop(float iouThres);
void rmLittle(int minW,int minH);
void rmBig(int maxW,int maxH);
void rmBigRatio(int maxW,int maxH,int imgH, float ratio);
void rmBigWithRoadMask(cv::Mat& mask,std::vector<float>& roadNumsList,float carLeastWidth);
void rmRatio();
std::vector<ObjBox> getRect();
};
}
#endif