feat: ✨ added bilateral filter calculator
This commit is contained in:
Arsenii Shulikov
parent
206813e3b0
commit
0363fe28a0
|
|
@ -69,8 +69,52 @@ cc_library(
|
|||
|
||||
|
||||
cc_library(
|
||||
name = "smooth_face_calculator",
|
||||
srcs = ["smooth_face_calculator.cc"],
|
||||
name = "smooth_face_calculator1",
|
||||
srcs = ["smooth_face1_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:vector",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "bilateral_filter_calculator",
|
||||
srcs = ["bilateral_filter_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:vector",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "smooth_face_calculator2",
|
||||
srcs = ["smooth_face2_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
|
|
|
|||
281
mediapipe/calculators/beauty/bilateral_filter_calculator.cc
Normal file
281
mediapipe/calculators/beauty/bilateral_filter_calculator.cc
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
// Copyright 2019 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_options.pb.h"
|
||||
#include "mediapipe/framework/formats/image_format.pb.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/vector.h"
|
||||
|
||||
namespace mediapipe
|
||||
{
|
||||
namespace
|
||||
{
|
||||
|
||||
constexpr char kMaskTag[] = "MASK";
|
||||
constexpr char kFaceBoxTag[] = "FACEBOX";
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kImageNewTag[] = "IMAGE2";
|
||||
|
||||
enum
|
||||
{
|
||||
ATTRIB_VERTEX,
|
||||
ATTRIB_TEXTURE_POSITION,
|
||||
NUM_ATTRIBUTES
|
||||
};
|
||||
|
||||
inline bool HasImageTag(mediapipe::CalculatorContext *cc) { return false; }
|
||||
} // namespace
|
||||
|
||||
class BilateralCalculator : public CalculatorBase
|
||||
{
|
||||
public:
|
||||
BilateralCalculator() = default;
|
||||
~BilateralCalculator() override = default;
|
||||
|
||||
static absl::Status GetContract(CalculatorContract *cc);
|
||||
|
||||
// From Calculator.
|
||||
absl::Status Open(CalculatorContext *cc) override;
|
||||
absl::Status Process(CalculatorContext *cc) override;
|
||||
absl::Status Close(CalculatorContext *cc) override;
|
||||
|
||||
private:
|
||||
absl::Status CreateRenderTargetCpu(CalculatorContext *cc,
|
||||
std::unique_ptr<cv::Mat> &image_mat,
|
||||
ImageFormat::Format *target_format);
|
||||
|
||||
absl::Status RenderToCpu(
|
||||
CalculatorContext *cc, const ImageFormat::Format &target_format,
|
||||
uchar *data_image);
|
||||
|
||||
absl::Status BilateralFilter(CalculatorContext *cc,
|
||||
const std::vector<double> &face_box);
|
||||
|
||||
// Indicates if image frame is available as input.
|
||||
bool image_frame_available_ = false;
|
||||
|
||||
int image_width_;
|
||||
int image_height_;
|
||||
cv::Mat mat_image_;
|
||||
std::unique_ptr<cv::Mat> image_mat;
|
||||
};
|
||||
REGISTER_CALCULATOR(BilateralCalculator);
|
||||
|
||||
absl::Status BilateralCalculator::GetContract(CalculatorContract *cc)
|
||||
{
|
||||
CHECK_GE(cc->Inputs().NumEntries(), 1);
|
||||
|
||||
if (cc->Inputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
CHECK(cc->Outputs().HasTag(kImageNewTag));
|
||||
}
|
||||
|
||||
// Data streams to render.
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id)
|
||||
{
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (tag == kFaceBoxTag)
|
||||
{
|
||||
cc->Inputs().Get(id).Set<std::vector<double>>();
|
||||
}
|
||||
else if (tag.empty())
|
||||
{
|
||||
// Empty tag defaults to accepting a single object of Mat type.
|
||||
cc->Inputs().Get(id).Set<cv::Mat>();
|
||||
}
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag(kImageNewTag))
|
||||
{
|
||||
cc->Outputs().Tag(kImageNewTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::Open(CalculatorContext *cc)
|
||||
{
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) || HasImageTag(cc))
|
||||
{
|
||||
image_frame_available_ = true;
|
||||
}
|
||||
|
||||
// Set the output header based on the input header (if present).
|
||||
const char *tag = kImageFrameTag;
|
||||
const char *out_tag = kImageNewTag;
|
||||
if (image_frame_available_ && !cc->Inputs().Tag(tag).Header().IsEmpty())
|
||||
{
|
||||
const auto &input_header =
|
||||
cc->Inputs().Tag(tag).Header().Get<VideoHeader>();
|
||||
auto *output_video_header = new VideoHeader(input_header);
|
||||
cc->Outputs().Tag(out_tag).SetHeader(Adopt(output_video_header));
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::Process(CalculatorContext *cc)
|
||||
{
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) &&
|
||||
cc->Inputs().Tag(kImageFrameTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kFaceBoxTag) &&
|
||||
cc->Inputs().Tag(kFaceBoxTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
// Initialize render target, drawn with OpenCV.
|
||||
ImageFormat::Format target_format;
|
||||
|
||||
if (cc->Outputs().HasTag(kImageNewTag))
|
||||
{
|
||||
MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, &target_format));
|
||||
}
|
||||
|
||||
mat_image_ = *image_mat.get();
|
||||
image_width_ = image_mat->cols;
|
||||
image_height_ = image_mat->rows;
|
||||
|
||||
const std::vector<double> &face_box =
|
||||
cc->Inputs().Tag(kFaceBoxTag).Get<std::vector<double>>();
|
||||
|
||||
MP_RETURN_IF_ERROR(BilateralFilter(cc, face_box));
|
||||
|
||||
// Copy the rendered image to output.
|
||||
uchar *image_mat_ptr = image_mat->data;
|
||||
MP_RETURN_IF_ERROR(RenderToCpu(cc, target_format, image_mat_ptr));
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::Close(CalculatorContext *cc)
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::RenderToCpu(
|
||||
CalculatorContext *cc, const ImageFormat::Format &target_format,
|
||||
uchar *data_image)
|
||||
{
|
||||
auto output_frame = absl::make_unique<ImageFrame>(
|
||||
target_format, image_width_, image_height_);
|
||||
|
||||
output_frame->CopyPixelData(target_format, image_width_, image_height_, data_image,
|
||||
ImageFrame::kDefaultAlignmentBoundary);
|
||||
|
||||
if (cc->Outputs().HasTag(kImageNewTag))
|
||||
{
|
||||
cc->Outputs()
|
||||
.Tag(kImageNewTag)
|
||||
.Add(output_frame.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::CreateRenderTargetCpu(
|
||||
CalculatorContext *cc, std::unique_ptr<cv::Mat> &image_mat,
|
||||
ImageFormat::Format *target_format)
|
||||
{
|
||||
if (image_frame_available_)
|
||||
{
|
||||
const auto &input_frame =
|
||||
cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
|
||||
int target_mat_type;
|
||||
switch (input_frame.Format())
|
||||
{
|
||||
case ImageFormat::SRGBA:
|
||||
*target_format = ImageFormat::SRGBA;
|
||||
target_mat_type = CV_8UC4;
|
||||
break;
|
||||
case ImageFormat::SRGB:
|
||||
*target_format = ImageFormat::SRGB;
|
||||
target_mat_type = CV_8UC3;
|
||||
break;
|
||||
case ImageFormat::GRAY8:
|
||||
*target_format = ImageFormat::SRGB;
|
||||
target_mat_type = CV_8UC3;
|
||||
break;
|
||||
default:
|
||||
return absl::UnknownError("Unexpected image frame format.");
|
||||
break;
|
||||
}
|
||||
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
input_frame.Height(), input_frame.Width(), target_mat_type);
|
||||
|
||||
auto input_mat = formats::MatView(&input_frame);
|
||||
|
||||
if (input_frame.Format() == ImageFormat::GRAY8)
|
||||
{
|
||||
cv::Mat rgb_mat;
|
||||
cv::cvtColor(input_mat, rgb_mat, CV_GRAY2RGB);
|
||||
rgb_mat.copyTo(*image_mat);
|
||||
}
|
||||
else
|
||||
{
|
||||
input_mat.copyTo(*image_mat);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
150, 150, CV_8UC4,
|
||||
cv::Scalar::all(255));
|
||||
*target_format = ImageFormat::SRGBA;
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status BilateralCalculator::BilateralFilter(CalculatorContext *cc,
|
||||
const std::vector<double> &face_box)
|
||||
{
|
||||
cv::Mat patch_face = mat_image_(cv::Range(face_box[1], face_box[3]),
|
||||
cv::Range(face_box[0], face_box[2]));
|
||||
cv::Mat patch_new, patch_wow;
|
||||
cv::cvtColor(patch_face, patch_wow, cv::COLOR_RGBA2RGB);
|
||||
cv::bilateralFilter(patch_wow, patch_new, 12, 50, 50);
|
||||
cv::cvtColor(patch_new, patch_new, cv::COLOR_RGB2RGBA);
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
|
@ -17,6 +17,7 @@
|
|||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <tuple>
|
||||
|
||||
#include <memory>
|
||||
|
||||
|
|
@ -53,11 +54,11 @@ namespace mediapipe
|
|||
inline bool HasImageTag(mediapipe::CalculatorContext *cc) { return false; }
|
||||
} // namespace
|
||||
|
||||
class SmoothFaceCalculator : public CalculatorBase
|
||||
class SmoothFaceCalculator1 : public CalculatorBase
|
||||
{
|
||||
public:
|
||||
SmoothFaceCalculator() = default;
|
||||
~SmoothFaceCalculator() override = default;
|
||||
SmoothFaceCalculator1() = default;
|
||||
~SmoothFaceCalculator1() override = default;
|
||||
|
||||
static absl::Status GetContract(CalculatorContract *cc);
|
||||
|
||||
|
|
@ -76,7 +77,6 @@ namespace mediapipe
|
|||
uchar *data_image);
|
||||
|
||||
absl::Status SmoothFace(CalculatorContext *cc,
|
||||
ImageFormat::Format *target_format,
|
||||
const std::unordered_map<std::string, cv::Mat> &mask_vec,
|
||||
const std::tuple<double, double, double, double> &face_box);
|
||||
|
||||
|
|
@ -88,11 +88,13 @@ namespace mediapipe
|
|||
int image_width_;
|
||||
int image_height_;
|
||||
cv::Mat mat_image_;
|
||||
cv::Mat not_full_face;
|
||||
std::vector<double> face_box;
|
||||
std::unique_ptr<cv::Mat> image_mat;
|
||||
};
|
||||
REGISTER_CALCULATOR(SmoothFaceCalculator);
|
||||
REGISTER_CALCULATOR(SmoothFaceCalculator1);
|
||||
|
||||
absl::Status SmoothFaceCalculator::GetContract(CalculatorContract *cc)
|
||||
absl::Status SmoothFaceCalculator1::GetContract(CalculatorContract *cc)
|
||||
{
|
||||
CHECK_GE(cc->Inputs().NumEntries(), 1);
|
||||
|
||||
|
|
@ -100,6 +102,7 @@ namespace mediapipe
|
|||
{
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
CHECK(cc->Outputs().HasTag(kImageFrameTag));
|
||||
CHECK(cc->Outputs().HasTag(kMaskTag));
|
||||
}
|
||||
|
||||
// Data streams to render.
|
||||
|
|
@ -128,11 +131,19 @@ namespace mediapipe
|
|||
{
|
||||
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
if (cc->Outputs().HasTag(kMaskTag))
|
||||
{
|
||||
cc->Outputs().Tag(kMaskTag).Set<cv::Mat>();
|
||||
}
|
||||
if (cc->Outputs().HasTag(kFaceBoxTag))
|
||||
{
|
||||
cc->Outputs().Tag(kFaceBoxTag).Set<std::vector<double>>();
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::Open(CalculatorContext *cc)
|
||||
absl::Status SmoothFaceCalculator1::Open(CalculatorContext *cc)
|
||||
{
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
|
|
@ -140,9 +151,6 @@ namespace mediapipe
|
|||
{
|
||||
image_frame_available_ = true;
|
||||
}
|
||||
else
|
||||
{
|
||||
}
|
||||
|
||||
// Set the output header based on the input header (if present).
|
||||
const char *tag = kImageFrameTag;
|
||||
|
|
@ -157,7 +165,7 @@ namespace mediapipe
|
|||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::Process(CalculatorContext *cc)
|
||||
absl::Status SmoothFaceCalculator1::Process(CalculatorContext *cc)
|
||||
{
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) &&
|
||||
cc->Inputs().Tag(kImageFrameTag).IsEmpty())
|
||||
|
|
@ -190,13 +198,13 @@ namespace mediapipe
|
|||
const std::vector<std::unordered_map<std::string, cv::Mat>> &mask_vec =
|
||||
cc->Inputs().Tag(kMaskTag).Get<std::vector<std::unordered_map<std::string, cv::Mat>>>();
|
||||
|
||||
const std::vector<std::tuple<double, double, double, double>> &face_box =
|
||||
const std::vector<std::tuple<double, double, double, double>> &face_boxes =
|
||||
cc->Inputs().Tag(kFaceBoxTag).Get<std::vector<std::tuple<double, double, double, double>>>();
|
||||
|
||||
if (mask_vec.size() > 0 && face_box.size() > 0)
|
||||
if (mask_vec.size() > 0 && face_boxes.size() > 0)
|
||||
{
|
||||
for (int i = 0; i < mask_vec.size(); i++)
|
||||
MP_RETURN_IF_ERROR(SmoothFace(cc, &target_format, mask_vec[i], face_box[i]));
|
||||
MP_RETURN_IF_ERROR(SmoothFace(cc, mask_vec[i], face_boxes[i]));
|
||||
}
|
||||
|
||||
// Copy the rendered image to output.
|
||||
|
|
@ -206,32 +214,50 @@ namespace mediapipe
|
|||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::Close(CalculatorContext *cc)
|
||||
absl::Status SmoothFaceCalculator1::Close(CalculatorContext *cc)
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::RenderToCpu(
|
||||
absl::Status SmoothFaceCalculator1::RenderToCpu(
|
||||
CalculatorContext *cc, const ImageFormat::Format &target_format,
|
||||
uchar *data_image)
|
||||
{
|
||||
auto output_frame = absl::make_unique<ImageFrame>(
|
||||
auto output_frame1 = absl::make_unique<ImageFrame>(
|
||||
target_format, image_width_, image_height_);
|
||||
|
||||
output_frame->CopyPixelData(target_format, image_width_, image_height_, data_image,
|
||||
output_frame1->CopyPixelData(target_format, image_width_, image_height_, data_image,
|
||||
ImageFrame::kDefaultAlignmentBoundary);
|
||||
|
||||
if (cc->Outputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
cc->Outputs()
|
||||
.Tag(kImageFrameTag)
|
||||
.Add(output_frame.release(), cc->InputTimestamp());
|
||||
.Add(output_frame1.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
auto output_frame2 = absl::make_unique<cv::Mat>(not_full_face);
|
||||
|
||||
if (cc->Outputs().HasTag(kMaskTag))
|
||||
{
|
||||
cc->Outputs()
|
||||
.Tag(kMaskTag)
|
||||
.Add(output_frame2.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
auto output_frame3 = absl::make_unique<std::vector<double>>(face_box);
|
||||
|
||||
if (cc->Outputs().HasTag(kFaceBoxTag))
|
||||
{
|
||||
cc->Outputs()
|
||||
.Tag(kFaceBoxTag)
|
||||
.Add(output_frame3.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::CreateRenderTargetCpu(
|
||||
absl::Status SmoothFaceCalculator1::CreateRenderTargetCpu(
|
||||
CalculatorContext *cc, std::unique_ptr<cv::Mat> &image_mat,
|
||||
ImageFormat::Format *target_format)
|
||||
{
|
||||
|
|
@ -287,7 +313,7 @@ namespace mediapipe
|
|||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
cv::Mat SmoothFaceCalculator::predict_forehead_mask(const std::unordered_map<std::string, cv::Mat> &mask_vec, double face_box_min_y)
|
||||
cv::Mat SmoothFaceCalculator1::predict_forehead_mask(const std::unordered_map<std::string, cv::Mat> &mask_vec, double face_box_min_y)
|
||||
{
|
||||
cv::Mat part_forehead_mask = mask_vec.find("PART_FOREHEAD_B")->second.clone();
|
||||
part_forehead_mask.convertTo(part_forehead_mask, CV_32F, 1.0 / 255);
|
||||
|
|
@ -357,20 +383,17 @@ namespace mediapipe
|
|||
return new_skin_mask;
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator::SmoothFace(CalculatorContext *cc,
|
||||
ImageFormat::Format *target_format,
|
||||
absl::Status SmoothFaceCalculator1::SmoothFace(CalculatorContext *cc,
|
||||
const std::unordered_map<std::string, cv::Mat> &mask_vec,
|
||||
const std::tuple<double, double, double, double> &face_box)
|
||||
const std::tuple<double, double, double, double> &face_boxx)
|
||||
{
|
||||
cv::Mat mouth_mask, mouth;
|
||||
|
||||
cv::Mat not_full_face = mask_vec.find("FACE_OVAL")->second.clone() -
|
||||
// predict_forehead_mask(mask_vec, std::get<1>(face_box)) -
|
||||
mask_vec.find("LEFT_EYE")->second.clone() -
|
||||
mask_vec.find("RIGHT_EYE")->second.clone() -
|
||||
mask_vec.find("LEFT_BROW")->second.clone() -
|
||||
mask_vec.find("RIGHT_BROW")->second.clone() -
|
||||
mask_vec.find("LIPS")->second.clone();
|
||||
not_full_face = mask_vec.find("FACE_OVAL")->second.clone() -
|
||||
// predict_forehead_mask(mask_vec, std::get<1>(face_boxx)) -
|
||||
mask_vec.find("LEFT_EYE")->second.clone() -
|
||||
mask_vec.find("RIGHT_EYE")->second.clone() -
|
||||
mask_vec.find("LEFT_BROW")->second.clone() -
|
||||
mask_vec.find("RIGHT_BROW")->second.clone() -
|
||||
mask_vec.find("LIPS")->second.clone();
|
||||
|
||||
cv::resize(not_full_face,
|
||||
not_full_face,
|
||||
|
|
@ -392,28 +415,10 @@ namespace mediapipe
|
|||
|
||||
cv::minMaxLoc(x, &min_x, &max_x);
|
||||
cv::minMaxLoc(y, &min_y, &max_y);
|
||||
|
||||
cv::Mat patch_face = mat_image_(cv::Range(min_y, max_y), cv::Range(min_x, max_x));
|
||||
cv::Mat patch_nff = not_full_face(cv::Range(min_y, max_y), cv::Range(min_x, max_x));
|
||||
cv::Mat patch_new, patch_wow;
|
||||
cv::cvtColor(patch_face, patch_wow, cv::COLOR_RGBA2RGB);
|
||||
cv::bilateralFilter(patch_wow, patch_new, 12, 50, 50);
|
||||
//patch_wow.copyTo(patch_new);
|
||||
|
||||
cv::Mat patch_new_nff, patch_new_mask, patch, patch_face_nff;
|
||||
|
||||
patch_new.copyTo(patch_new_nff, patch_nff);
|
||||
|
||||
patch_face.copyTo(patch_face_nff, patch_nff);
|
||||
cv::cvtColor(patch_face_nff, patch_face_nff, cv::COLOR_RGBA2RGB);
|
||||
|
||||
patch_new_mask = 0.85 * patch_new_nff + 0.15 * patch_face_nff;
|
||||
|
||||
patch = cv::min(255, patch_new_mask);
|
||||
|
||||
cv::cvtColor(patch, patch, cv::COLOR_RGB2RGBA);
|
||||
|
||||
patch.copyTo(patch_face, patch_nff);
|
||||
face_box.push_back(min_x);
|
||||
face_box.push_back(min_y);
|
||||
face_box.push_back(max_x);
|
||||
face_box.push_back(max_y);
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
342
mediapipe/calculators/beauty/smooth_face2_calculator.cc
Normal file
342
mediapipe/calculators/beauty/smooth_face2_calculator.cc
Normal file
|
|
@ -0,0 +1,342 @@
|
|||
// Copyright 2019 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <tuple>
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_options.pb.h"
|
||||
#include "mediapipe/framework/formats/image_format.pb.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/vector.h"
|
||||
|
||||
namespace mediapipe
|
||||
{
|
||||
namespace
|
||||
{
|
||||
|
||||
constexpr char kMaskTag[] = "MASK";
|
||||
constexpr char kFaceBoxTag[] = "FACEBOX";
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kImageNewTag[] = "IMAGE2";
|
||||
|
||||
enum
|
||||
{
|
||||
ATTRIB_VERTEX,
|
||||
ATTRIB_TEXTURE_POSITION,
|
||||
NUM_ATTRIBUTES
|
||||
};
|
||||
|
||||
inline bool HasImageTag(mediapipe::CalculatorContext *cc) { return false; }
|
||||
} // namespace
|
||||
|
||||
class SmoothFaceCalculator2 : public CalculatorBase
|
||||
{
|
||||
public:
|
||||
SmoothFaceCalculator2() = default;
|
||||
~SmoothFaceCalculator2() override = default;
|
||||
|
||||
static absl::Status GetContract(CalculatorContract *cc);
|
||||
|
||||
// From Calculator.
|
||||
absl::Status Open(CalculatorContext *cc) override;
|
||||
absl::Status Process(CalculatorContext *cc) override;
|
||||
absl::Status Close(CalculatorContext *cc) override;
|
||||
|
||||
private:
|
||||
absl::Status CreateRenderTargetCpu(CalculatorContext *cc,
|
||||
std::unique_ptr<cv::Mat> &image_mat,
|
||||
std::unique_ptr<cv::Mat> &new_mat,
|
||||
ImageFormat::Format *target_format);
|
||||
|
||||
absl::Status RenderToCpu(
|
||||
CalculatorContext *cc, const ImageFormat::Format &target_format,
|
||||
uchar *data_image);
|
||||
|
||||
absl::Status SmoothEnd(CalculatorContext *cc,
|
||||
const std::vector<double> &face_box);
|
||||
|
||||
// Indicates if image frame is available as input.
|
||||
bool image_frame_available_ = false;
|
||||
|
||||
int image_width_;
|
||||
int image_height_;
|
||||
cv::Mat mat_image_;
|
||||
cv::Mat new_image_;
|
||||
cv::Mat not_full_face;
|
||||
std::unique_ptr<cv::Mat> image_mat;
|
||||
std::unique_ptr<cv::Mat> new_mat;
|
||||
std::unique_ptr<cv::Mat> nff_mat;
|
||||
};
|
||||
REGISTER_CALCULATOR(SmoothFaceCalculator2);
|
||||
|
||||
absl::Status SmoothFaceCalculator2::GetContract(CalculatorContract *cc)
|
||||
{
|
||||
CHECK_GE(cc->Inputs().NumEntries(), 1);
|
||||
|
||||
if (cc->Inputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
CHECK(cc->Outputs().HasTag(kImageFrameTag));
|
||||
}
|
||||
if (cc->Inputs().HasTag(kImageNewTag))
|
||||
{
|
||||
cc->Inputs().Tag(kImageNewTag).Set<ImageFrame>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kMaskTag))
|
||||
{
|
||||
cc->Inputs().Tag(kMaskTag).Set<cv::Mat>();
|
||||
}
|
||||
|
||||
// Data streams to render.
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id)
|
||||
{
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (tag == kFaceBoxTag)
|
||||
{
|
||||
cc->Inputs().Get(id).Set<std::vector<double>>();
|
||||
}
|
||||
else if (tag.empty())
|
||||
{
|
||||
// Empty tag defaults to accepting a single object of Mat type.
|
||||
cc->Inputs().Get(id).Set<cv::Mat>();
|
||||
}
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator2::Open(CalculatorContext *cc)
|
||||
{
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) || HasImageTag(cc))
|
||||
{
|
||||
image_frame_available_ = true;
|
||||
}
|
||||
|
||||
// Set the output header based on the input header (if present).
|
||||
const char *tag = kImageFrameTag;
|
||||
if (image_frame_available_ && !cc->Inputs().Tag(tag).Header().IsEmpty())
|
||||
{
|
||||
const auto &input_header =
|
||||
cc->Inputs().Tag(tag).Header().Get<VideoHeader>();
|
||||
auto *output_video_header = new VideoHeader(input_header);
|
||||
cc->Outputs().Tag(tag).SetHeader(Adopt(output_video_header));
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator2::Process(CalculatorContext *cc)
|
||||
{
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) &&
|
||||
cc->Inputs().Tag(kImageFrameTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kImageNewTag) &&
|
||||
cc->Inputs().Tag(kImageNewTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kMaskTag) &&
|
||||
cc->Inputs().Tag(kMaskTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kFaceBoxTag) &&
|
||||
cc->Inputs().Tag(kFaceBoxTag).IsEmpty())
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
// Initialize render target, drawn with OpenCV.
|
||||
ImageFormat::Format target_format;
|
||||
|
||||
if (cc->Outputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, new_mat, &target_format));
|
||||
}
|
||||
|
||||
not_full_face = cc->Inputs().Tag(kMaskTag).Get<cv::Mat>();
|
||||
new_image_ = *new_mat.get();
|
||||
mat_image_ = *image_mat.get();
|
||||
image_width_ = image_mat->cols;
|
||||
image_height_ = image_mat->rows;
|
||||
|
||||
const std::vector<double> &face_box =
|
||||
cc->Inputs().Tag(kFaceBoxTag).Get<std::vector<double>>();
|
||||
|
||||
MP_RETURN_IF_ERROR(SmoothEnd(cc, face_box));
|
||||
|
||||
// Copy the rendered image to output.
|
||||
uchar *image_mat_ptr = image_mat->data;
|
||||
MP_RETURN_IF_ERROR(RenderToCpu(cc, target_format, image_mat_ptr));
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator2::Close(CalculatorContext *cc)
|
||||
{
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator2::RenderToCpu(
|
||||
CalculatorContext *cc, const ImageFormat::Format &target_format,
|
||||
uchar *data_image)
|
||||
{
|
||||
auto output_frame = absl::make_unique<ImageFrame>(
|
||||
target_format, image_width_, image_height_);
|
||||
|
||||
output_frame->CopyPixelData(target_format, image_width_, image_height_, data_image,
|
||||
ImageFrame::kDefaultAlignmentBoundary);
|
||||
|
||||
if (cc->Outputs().HasTag(kImageFrameTag))
|
||||
{
|
||||
cc->Outputs()
|
||||
.Tag(kImageFrameTag)
|
||||
.Add(output_frame.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SmoothFaceCalculator2::CreateRenderTargetCpu(
|
||||
CalculatorContext *cc,
|
||||
std::unique_ptr<cv::Mat> &image_mat,
|
||||
std::unique_ptr<cv::Mat> &new_mat,
|
||||
ImageFormat::Format *target_format)
|
||||
{
|
||||
if (image_frame_available_)
|
||||
{
|
||||
const auto &input_frame =
|
||||
cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
const auto &new_frame =
|
||||
cc->Inputs().Tag(kImageNewTag).Get<ImageFrame>();
|
||||
|
||||
int target_mat_type;
|
||||
switch (input_frame.Format())
|
||||
{
|
||||
case ImageFormat::SRGBA:
|
||||
*target_format = ImageFormat::SRGBA;
|
||||
target_mat_type = CV_8UC4;
|
||||
break;
|
||||
case ImageFormat::SRGB:
|
||||
*target_format = ImageFormat::SRGB;
|
||||
target_mat_type = CV_8UC3;
|
||||
break;
|
||||
case ImageFormat::GRAY8:
|
||||
*target_format = ImageFormat::SRGB;
|
||||
target_mat_type = CV_8UC3;
|
||||
break;
|
||||
default:
|
||||
return absl::UnknownError("Unexpected image frame format.");
|
||||
break;
|
||||
}
|
||||
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
input_frame.Height(), input_frame.Width(), target_mat_type);
|
||||
|
||||
auto input_mat = formats::MatView(&input_frame);
|
||||
|
||||
new_mat = absl::make_unique<cv::Mat>(
|
||||
new_frame.Height(), new_frame.Width(), target_mat_type);
|
||||
|
||||
auto new_input_mat = formats::MatView(&new_frame);
|
||||
|
||||
if (input_frame.Format() == ImageFormat::GRAY8)
|
||||
{
|
||||
cv::Mat rgb_mat;
|
||||
cv::Mat rgb_mat2;
|
||||
cv::cvtColor(input_mat, rgb_mat, CV_GRAY2RGB);
|
||||
rgb_mat.copyTo(*image_mat);
|
||||
cv::cvtColor(new_input_mat, rgb_mat2, CV_GRAY2RGB);
|
||||
rgb_mat2.copyTo(*new_mat);
|
||||
}
|
||||
else
|
||||
{
|
||||
input_mat.copyTo(*image_mat);
|
||||
new_input_mat.copyTo(*new_mat);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
150, 150, CV_8UC4,
|
||||
cv::Scalar::all(255));
|
||||
nff_mat = absl::make_unique<cv::Mat>(
|
||||
150, 150, CV_8UC4,
|
||||
cv::Scalar::all(255));
|
||||
new_mat = absl::make_unique<cv::Mat>(
|
||||
150, 150, CV_8UC4,
|
||||
cv::Scalar::all(255));
|
||||
*target_format = ImageFormat::SRGBA;
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
absl::Status SmoothFaceCalculator2::SmoothEnd(CalculatorContext *cc,
|
||||
const std::vector<double> &face_box)
|
||||
{
|
||||
cv::Mat patch_face = mat_image_(cv::Range(face_box[1], face_box[3]),
|
||||
cv::Range(face_box[0], face_box[2]));
|
||||
cv::Mat patch_new = new_image_(cv::Range(face_box[1], face_box[3]),
|
||||
cv::Range(face_box[0], face_box[2]));
|
||||
cv::Mat patch_nff = not_full_face(cv::Range(face_box[1], face_box[3]),
|
||||
cv::Range(face_box[0], face_box[2]));
|
||||
|
||||
cv::Mat patch_new_nff, patch_new_mask, patch, patch_face_nff;
|
||||
|
||||
patch_new.copyTo(patch_new_nff, patch_nff);
|
||||
|
||||
patch_face.copyTo(patch_face_nff, patch_nff);
|
||||
cv::cvtColor(patch_face_nff, patch_face_nff, cv::COLOR_RGBA2RGB);
|
||||
cv::cvtColor(patch_new_nff, patch_new_nff, cv::COLOR_RGBA2RGB);
|
||||
|
||||
patch_new_mask = 0.85 * patch_new_nff + 0.15 * patch_face_nff;
|
||||
|
||||
patch = cv::min(255, patch_new_mask);
|
||||
|
||||
cv::cvtColor(patch, patch, cv::COLOR_RGB2RGBA);
|
||||
|
||||
patch.copyTo(patch_face, patch_nff);
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
|
@ -27,7 +27,9 @@ cc_library(
|
|||
"//mediapipe/calculators/core:split_proto_list_calculator",
|
||||
"//mediapipe/util:annotation_renderer",
|
||||
"//mediapipe/calculators/util:annotation_overlay_calculator",
|
||||
"//mediapipe/calculators/beauty:smooth_face_calculator",
|
||||
"//mediapipe/calculators/beauty:smooth_face_calculator1",
|
||||
"//mediapipe/calculators/beauty:bilateral_filter_calculator",
|
||||
"//mediapipe/calculators/beauty:smooth_face_calculator2",
|
||||
"//mediapipe/calculators/beauty:draw_lipstick_calculator",
|
||||
"//mediapipe/calculators/beauty:whiten_teeth_calculator",
|
||||
"//mediapipe/calculators/util:detections_to_render_data_calculator",
|
||||
|
|
|
|||
|
|
@ -74,9 +74,29 @@ node {
|
|||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "SmoothFaceCalculator"
|
||||
calculator: "SmoothFaceCalculator1"
|
||||
input_stream: "IMAGE:input_image_2"
|
||||
input_stream: "MASK:0:multi_mask"
|
||||
input_stream: "FACEBOX:0:multi_face_box"
|
||||
output_stream: "IMAGE:input_image_3"
|
||||
output_stream: "MASK:not_full_face"
|
||||
output_stream: "FACEBOX:box1"
|
||||
}
|
||||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "BilateralCalculator"
|
||||
input_stream: "IMAGE:input_image_3"
|
||||
input_stream: "FACEBOX:box1"
|
||||
output_stream: "IMAGE2:input_image_4"
|
||||
}
|
||||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "SmoothFaceCalculator2"
|
||||
input_stream: "IMAGE:input_image_2"
|
||||
input_stream: "IMAGE2:input_image_4"
|
||||
input_stream: "MASK:not_full_face"
|
||||
input_stream: "FACEBOX:box1"
|
||||
output_stream: "IMAGE:output_image"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -78,9 +78,29 @@ node {
|
|||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "SmoothFaceCalculator"
|
||||
input_stream: "IMAGE_GPU:input_image_2"
|
||||
calculator: "SmoothFaceCalculator1"
|
||||
input_stream: "IMAGE:input_image_2"
|
||||
input_stream: "MASK:0:multi_mask"
|
||||
input_stream: "FACEBOX:0:multi_face_box"
|
||||
output_stream: "IMAGE_GPU:output_image"
|
||||
output_stream: "IMAGE:input_image_3"
|
||||
output_stream: "MASK:not_full_face"
|
||||
output_stream: "FACEBOX:box1"
|
||||
}
|
||||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "BilateralCalculator"
|
||||
input_stream: "IMAGE:input_image_3"
|
||||
input_stream: "FACEBOX:box1"
|
||||
output_stream: "IMAGE2:input_image_4"
|
||||
}
|
||||
|
||||
#Smoothes face on the IMAGE using MASK.
|
||||
node {
|
||||
calculator: "SmoothFaceCalculator2"
|
||||
input_stream: "IMAGE:input_image_2"
|
||||
input_stream: "IMAGE2:input_image_4"
|
||||
input_stream: "MASK:not_full_face"
|
||||
input_stream: "FACEBOX:box1"
|
||||
output_stream: "IMAGE:output_image"
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user