diff --git a/fastdeploy/vision/visualize/matting.cc b/fastdeploy/vision/visualize/matting.cc index abc245e73..3dd34bffa 100644 --- a/fastdeploy/vision/visualize/matting.cc +++ b/fastdeploy/vision/visualize/matting.cc @@ -19,8 +19,8 @@ namespace fastdeploy { namespace vision { cv::Mat VisMatting(const cv::Mat& im, const MattingResult& result, - bool remove_small_connected_area, bool is_transparent, - float transparent_threshold) { + bool transparent_background, float transparent_threshold, + bool remove_small_connected_area) { FDASSERT((!im.empty()), "im can't be empty!"); auto vis_img = im.clone(); @@ -44,8 +44,8 @@ cv::Mat VisMatting(const cv::Mat& im, const MattingResult& result, if ((vis_img).type() != CV_8UC3) { (vis_img).convertTo((vis_img), CV_8UC3); } - - if(is_transparent) { + + if (transparent_background) { if (vis_img.channels() != 4) { cv::cvtColor(vis_img, transparent_vis_mat, cv::COLOR_BGR2BGRA); vis_img = transparent_vis_mat; @@ -60,17 +60,24 @@ cv::Mat VisMatting(const cv::Mat& im, const MattingResult& result, for (size_t i = 0; i < height; ++i) { for (size_t j = 0; j < width; ++j) { float alpha_val = alpha_data[i * width + j]; - vis_data[i * width * channel + j * channel + 0] = cv::saturate_cast( - static_cast(im_data[i * width * 3 + j * 3 + 0]) * alpha_val + - (1.f - alpha_val) * 153.f); - vis_data[i * width * channel + j * channel + 1] = cv::saturate_cast( - static_cast(im_data[i * width * 3 + j * 3 + 1]) * alpha_val + - (1.f - alpha_val) * 255.f); - vis_data[i * width * channel + j * channel + 2] = cv::saturate_cast( - static_cast(im_data[i * width * 3 + j * 3 + 2]) * alpha_val + - (1.f - alpha_val) * 120.f); - if (is_transparent && alpha_val < transparent_threshold) { - vis_data[i * width * channel + j * channel + 3] = cv::saturate_cast(0.f); + vis_data[i * width * channel + j * channel + 0] = + cv::saturate_cast( + static_cast(im_data[i * width * 3 + j * 3 + 0]) * + alpha_val + + (1.f - alpha_val) * 153.f); + vis_data[i * width * channel + j * channel + 1] = + cv::saturate_cast( + static_cast(im_data[i * width * 3 + j * 3 + 1]) * + alpha_val + + (1.f - alpha_val) * 255.f); + vis_data[i * width * channel + j * channel + 2] = + cv::saturate_cast( + static_cast(im_data[i * width * 3 + j * 3 + 2]) * + alpha_val + + (1.f - alpha_val) * 120.f); + if (transparent_background && alpha_val < transparent_threshold) { + vis_data[i * width * channel + j * channel + 3] = + cv::saturate_cast(0.f); } } } diff --git a/fastdeploy/vision/visualize/visualize.h b/fastdeploy/vision/visualize/visualize.h index 9f42ca083..c3a77e394 100755 --- a/fastdeploy/vision/visualize/visualize.h +++ b/fastdeploy/vision/visualize/visualize.h @@ -143,16 +143,16 @@ FASTDEPLOY_DECL cv::Mat VisSegmentation(const cv::Mat& im, * * \param[in] im the input image data, comes from cv::imread(), is a 3-D array with layout HWC, BGR format * \param[in] result the result produced by model - * \param[in] remove_small_connected_area if remove_small_connected_area==true, the visualized result will not include the small connected areas * \param[in] is_transparent if is_transparent==true, the background will with transparent color * \param[in] transparent_threshold since the alpha value in MattringResult is a float between [0, 1], transparent_threshold is used to filter background pixel + * \param[in] remove_small_connected_area if remove_small_connected_area==true, the visualized result will not include the small connected areas * \return cv::Mat type stores the visualized results */ FASTDEPLOY_DECL cv::Mat VisMatting(const cv::Mat& im, const MattingResult& result, - bool remove_small_connected_area = false, bool is_transparent = false, - float transparent_threshold = 0.999); + float transparent_threshold = 0.999 + bool remove_small_connected_area = false); /** \brief Show the visualized results for Ocr models * * \param[in] im the input image data, comes from cv::imread(), is a 3-D array with layout HWC, BGR format diff --git a/fastdeploy/vision/visualize/visualize_pybind.cc b/fastdeploy/vision/visualize/visualize_pybind.cc index 802790c30..c5b43f9d1 100644 --- a/fastdeploy/vision/visualize/visualize_pybind.cc +++ b/fastdeploy/vision/visualize/visualize_pybind.cc @@ -102,10 +102,12 @@ void BindVisualize(pybind11::module& m) { }) .def("vis_matting", [](pybind11::array& im_data, vision::MattingResult& result, + bool transparent_background, bool transparent_threshold, bool remove_small_connected_area) { cv::Mat im = PyArrayToCvMat(im_data); - auto vis_im = - vision::VisMatting(im, result, remove_small_connected_area); + auto vis_im = vision::VisMatting( + im, result, transparent_background, transparent_threshold, + remove_small_connected_area); FDTensor out; vision::Mat(vis_im).ShareWithTensor(&out); return TensorToPyArray(out); diff --git a/python/fastdeploy/vision/visualize/__init__.py b/python/fastdeploy/vision/visualize/__init__.py index e528c8804..b87a2dfda 100755 --- a/python/fastdeploy/vision/visualize/__init__.py +++ b/python/fastdeploy/vision/visualize/__init__.py @@ -95,15 +95,22 @@ def vis_matting_alpha(im_data, remove_small_connected_area) -def vis_matting(im_data, matting_result, remove_small_connected_area=False): +def vis_matting(im_data, + matting_result, + transparent_background=False, + transparent_threshold=0.99, + remove_small_connected_area=False): """Show the visualized results for matting models :param im_data: (numpy.ndarray)The input image data, 3-D array with layout HWC, BGR format :param matting_result: the result produced by model + :param transparent_background: whether visulizing matting result with transparent background + :param transparent_threshold: since the alpha value in MattringResult is a float between [0, 1], transparent_threshold is used to filter background pixel :param remove_small_connected_area: (bool) if remove_small_connected_area==True, the visualized result will not include the small connected areas :return: (numpy.ndarray) image with visualized results """ return C.vision.vis_matting(im_data, matting_result, + transparent_background, transparent_threshold, remove_small_connected_area)