From 6a74b88b4e20ca35bcae037f968ed25c0c7eea5e Mon Sep 17 00:00:00 2001 From: Huang ZhuangZhuang Date: Tue, 11 May 2021 19:57:11 +0800 Subject: [PATCH] update retinaface_trt del cvtColor and increase THRESH (#535) * update retinaface_trt del cvtColor and increase THRESH * update macro in retinafce_xx.cpp and preprocess_image description * retinaface_xx.cpp add a space after * HRESH for CONF and THRESHOLD for IOU is unified --THRESH Co-authored-by: Admin --- retinaface/retina_mnet.cpp | 7 ++++--- retinaface/retina_r50.cpp | 6 ++++-- retinaface/retinaface_trt.py | 12 +++++------- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/retinaface/retina_mnet.cpp b/retinaface/retina_mnet.cpp index 14c3ec7..8555df5 100644 --- a/retinaface/retina_mnet.cpp +++ b/retinaface/retina_mnet.cpp @@ -12,7 +12,8 @@ #define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32 #define DEVICE 0 // GPU id #define BATCH_SIZE 1 -#define VIS_THRESH 0.6 +#define CONF_THRESH = 0.75 +#define IOU_THRESH = 0.4 // stuff we know about the network and the input/output blobs static const int INPUT_H = decodeplugin::INPUT_H; // H, W must be able to be divided by 32. @@ -349,13 +350,13 @@ int main(int argc, char** argv) { for (int b = 0; b < BATCH_SIZE; b++) { std::vector res; - nms(res, &prob[b * OUTPUT_SIZE]); + nms(res, &prob[b * OUTPUT_SIZE], IOU_THRESH); std::cout << "number of detections -> " << prob[b * OUTPUT_SIZE] << std::endl; std::cout << " -> " << prob[b * OUTPUT_SIZE + 10] << std::endl; std::cout << "after nms -> " << res.size() << std::endl; cv::Mat tmp = img.clone(); for (size_t j = 0; j < res.size(); j++) { - if (res[j].class_confidence < VIS_THRESH) continue; + if (res[j].class_confidence < CONF_THRESH) continue; cv::Rect r = get_rect_adapt_landmark(tmp, INPUT_W, INPUT_H, res[j].bbox, res[j].landmark); cv::rectangle(tmp, r, cv::Scalar(0x27, 0xC1, 0x36), 2); //cv::putText(tmp, std::to_string((int)(res[j].class_confidence * 100)) + "%", cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 1); diff --git a/retinaface/retina_r50.cpp b/retinaface/retina_r50.cpp index f857f59..d2e51d1 100644 --- a/retinaface/retina_r50.cpp +++ b/retinaface/retina_r50.cpp @@ -12,6 +12,8 @@ #define USE_INT8 // set USE_INT8 or USE_FP16 or USE_FP32 #define DEVICE 0 // GPU id #define BATCH_SIZE 1 +#define CONF_THRESH = 0.75 +#define IOU_THRESH = 0.4 // stuff we know about the network and the input/output blobs static const int INPUT_H = decodeplugin::INPUT_H; // H, W must be able to be divided by 32. @@ -369,13 +371,13 @@ int main(int argc, char** argv) { for (int b = 0; b < BATCH_SIZE; b++) { std::vector res; - nms(res, &prob[b * OUTPUT_SIZE]); + nms(res, &prob[b * OUTPUT_SIZE], IOU_THRESH); std::cout << "number of detections -> " << prob[b * OUTPUT_SIZE] << std::endl; std::cout << " -> " << prob[b * OUTPUT_SIZE + 10] << std::endl; std::cout << "after nms -> " << res.size() << std::endl; cv::Mat tmp = img.clone(); for (size_t j = 0; j < res.size(); j++) { - if (res[j].class_confidence < 0.1) continue; + if (res[j].class_confidence < CONF_THRESH) continue; cv::Rect r = get_rect_adapt_landmark(tmp, INPUT_W, INPUT_H, res[j].bbox, res[j].landmark); cv::rectangle(tmp, r, cv::Scalar(0x27, 0xC1, 0x36), 2); //cv::putText(tmp, std::to_string((int)(res[j].class_confidence * 100)) + "%", cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 1); diff --git a/retinaface/retinaface_trt.py b/retinaface/retinaface_trt.py index 8a148cd..65a78c7 100644 --- a/retinaface/retinaface_trt.py +++ b/retinaface/retinaface_trt.py @@ -19,8 +19,8 @@ import torchvision INPUT_H = 480 #defined in decode.h INPUT_W = 640 -CONF_THRESH = 0.4 -IOU_THRESHOLD = 0.1 +CONF_THRESH = 0.75 +IOU_THRESHOLD = 0.4 np.set_printoptions(threshold=np.inf) def plot_one_box(x, landmark,img, color=None, label=None, line_thickness=None): @@ -180,9 +180,8 @@ class Retinaface_trt(object): def preprocess_image(self, input_image_path): """ - description: Read an image from image path, convert it to RGB, - resize and pad it to target size, normalize to [0,1], - transform to NCHW format. + description: Read an image from image path, resize and pad it to target size, + normalize to [0,1],transform to NCHW format. param: input_image_path: str, image path return: @@ -193,8 +192,7 @@ class Retinaface_trt(object): """ image_raw = cv2.imread(input_image_path) h, w, c = image_raw.shape - image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB) - image = cv2.resize(image, (INPUT_W, INPUT_H)) + image = cv2.resize(image_raw, (INPUT_W, INPUT_H)) image = image.astype(np.float32)