cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

WarpAffineCalculator on api2

Browse Source
This commit is contained in:
mslight 2022-08-08 15:43:25 +04:00
parent 51ed94633e
commit ffa2853b79
3 changed files with 105 additions and 249 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mediapipe/graphs/deformation/calculators/Tensor.h
View File

@ -64,10 +64,6 @@ public:
return vector<T>(M.ptr<T>(0), M.ptr<T>(0) + dims[1]); return vector<T>(M.ptr<T>(0), M.ptr<T>(0) + dims[1]);
} }
/* vector<int> get_dims() {
return dims;
} */
T at(vector<int> _indexes) { T at(vector<int> _indexes) {
return M.at<T>(_indexes.data()); return M.at<T>(_indexes.data());
} }

8
mediapipe/graphs/deformation/calculators/face_processor_calculator.cc
View File

@ -283,9 +283,9 @@ namespace mediapipe
absl::Status FaceProcessorCalculator::ProcessImage(CalculatorContext *cc) absl::Status FaceProcessorCalculator::ProcessImage(CalculatorContext *cc)
{ {
double alfaNose = 2.7; double alfaNose = 1.2;
double alfaLips = 0.7; double alfaLips = 0.4;
double alfaCheekbones = 0.7; double alfaCheekbones = 0.4;
if (cc->Inputs().HasTag(kNormLandmarksTag)) if (cc->Inputs().HasTag(kNormLandmarksTag))
{ {
@ -381,7 +381,7 @@ namespace mediapipe
Tensor<double> _src = __facePts.index(_trianglesIndexes).index(_order); Tensor<double> _src = __facePts.index(_trianglesIndexes).index(_order);
Tensor<double> _dst = ___facePts.index(_trianglesIndexes).index(_order); Tensor<double> _dst = ___facePts.index(_trianglesIndexes).index(_order);
// cout << _src.get_dims().size() << endl;
auto srcPtr = absl::make_unique<Tensor<double>>(_src); auto srcPtr = absl::make_unique<Tensor<double>>(_src);
cc->Outputs().Tag(kSrcTensorTag).Add(srcPtr.release(), cc->InputTimestamp()); cc->Outputs().Tag(kSrcTensorTag).Add(srcPtr.release(), cc->InputTimestamp());

216
mediapipe/graphs/deformation/calculators/warp_affine_calculator.cc
View File

@ -14,6 +14,7 @@
#include "Tensor.h" #include "Tensor.h"
#include "mediapipe/framework/calculator_framework.h" #include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/formats/image_format.pb.h" #include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h" #include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h" #include "mediapipe/framework/formats/image_frame_opencv.h"
@ -23,201 +24,48 @@
#include "mediapipe/framework/port/opencv_highgui_inc.h" #include "mediapipe/framework/port/opencv_highgui_inc.h"
#include "mediapipe/framework/port/status.h" #include "mediapipe/framework/port/status.h"
using namespace cv;
using namespace std;
namespace mediapipe namespace mediapipe
{ {
namespace namespace api2
{ {
constexpr char kImageFrameTag[] = "IMAGE"; class WarpAffineCalculator : public Node
constexpr char kSrcTag[] = "SRC_TENSOR";
constexpr char kDstTag[] = "DST_TENSOR";
inline bool HasImageTag(mediapipe::CalculatorContext *cc) { return false; }
} // namespace
class WarpAffineCalculator : public CalculatorBase
{ {
public: public:
WarpAffineCalculator() = default; static constexpr Input<ImageFrame> kImageFrame{"IMAGE"};
~WarpAffineCalculator() override = default; static constexpr Input<Tensor<double>>::Optional kSrc{"SRC_TENSOR"};
static constexpr Input<Tensor<double>>::Optional kDst{"DST_TENSOR"};
static constexpr Output<ImageFrame> kOut{"IMAGE"};
static absl::Status GetContract(CalculatorContract *cc); MEDIAPIPE_NODE_CONTRACT(kImageFrame, kSrc, kDst, kOut);
absl::Status Open(CalculatorContext *cc) override; static absl::Status UpdateContract(CalculatorContract *cc)
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, std::unique_ptr<cv::Mat> &image_mat);
absl::Status AffineTransform(CalculatorContext *cc, std::unique_ptr<cv::Mat> &image_mat, Tensor<double> _src, Tensor<double> _dst);
bool image_frame_available_ = false;
std::unique_ptr<Mat> image_mat;
};
absl::Status WarpAffineCalculator::GetContract(CalculatorContract *cc)
{ {
RET_CHECK(cc->Inputs().HasTag(kImageFrameTag)); RET_CHECK(kOut(cc).IsConnected())
<< "At least one output stream is expected.";
if (cc->Inputs().HasTag(kImageFrameTag))
{
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
}
if (cc->Inputs().HasTag(kSrcTag))
{
cc->Inputs().Tag(kSrcTag).Set<Tensor<double>>();
}
if (cc->Inputs().HasTag(kDstTag))
{
cc->Inputs().Tag(kDstTag).Set<Tensor<double>>();
}
if (cc->Outputs().HasTag(kImageFrameTag))
{
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
}
return absl::OkStatus(); return absl::OkStatus();
} }
absl::Status WarpAffineCalculator::Open(CalculatorContext *cc) absl::Status Process(CalculatorContext *cc) override
{ {
cc->SetOffset(TimestampDiff(0)); if (kImageFrame(cc).IsEmpty())
if (cc->Inputs().HasTag(kImageFrameTag) || HasImageTag(cc))
{
image_frame_available_ = true;
}
return absl::OkStatus();
}
absl::Status WarpAffineCalculator::Process(CalculatorContext *cc)
{
if (cc->Inputs().Tag(kImageFrameTag).IsEmpty())
{ {
return absl::OkStatus(); return absl::OkStatus();
} }
ImageFormat::Format target_format; const ImageFrame &input = *kImageFrame(cc);
auto image = absl::make_unique<cv::Mat>(
input.Height(), input.Width(), CV_8UC3);
MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, &target_format)); auto mat_image_ = formats::MatView(&input);
if (!cc->Inputs().Tag(kSrcTag).IsEmpty() && !cc->Inputs().Tag(kDstTag).IsEmpty())
{
const Tensor<double> _src = cc->Inputs().Tag(kSrcTag).Get<Tensor<double>>();
const Tensor<double> _dst = cc->Inputs().Tag(kDstTag).Get<Tensor<double>>();
MP_RETURN_IF_ERROR(AffineTransform(cc, image_mat, _src, _dst));
}
// Copy the rendered image to output.
uchar *image_mat_ptr = image_mat->data;
MP_RETURN_IF_ERROR(RenderToCpu(cc, target_format, image_mat_ptr, image_mat));
return absl::OkStatus();
}
absl::Status WarpAffineCalculator::Close(CalculatorContext *cc)
{
return absl::OkStatus();
}
absl::Status WarpAffineCalculator::RenderToCpu(
CalculatorContext *cc, const ImageFormat::Format &target_format,
uchar *data_image, std::unique_ptr<cv::Mat> &image_mat)
{
auto output_frame = absl::make_unique<ImageFrame>(
target_format, image_mat->cols, image_mat->rows);
output_frame->CopyPixelData(target_format, image_mat->cols, image_mat->rows, data_image,
ImageFrame::kDefaultAlignmentBoundary);
if (cc->Outputs().HasTag(kImageFrameTag))
{
cc->Outputs()
.Tag(kImageFrameTag)
.Add(output_frame.release(), cc->InputTimestamp());
}
return absl::OkStatus();
}
absl::Status WarpAffineCalculator::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::SBGR:
*target_format = ImageFormat::SBGR;
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>(
1920, 1080, CV_8UC4,
cv::Scalar(cv::Scalar::all(255)));
*target_format = ImageFormat::SRGBA;
}
return absl::OkStatus();
}
absl::Status WarpAffineCalculator::AffineTransform(CalculatorContext *cc, std::unique_ptr<cv::Mat> &image_mat, Tensor<double> _src, Tensor<double> _dst)
{
Mat mat_image_ = *image_mat.get();
Mat clone_image = mat_image_.clone();
Mat outImage = Mat(mat_image_.size(), mat_image_.type());
Mat out = mat_image_.clone(); Mat out = mat_image_.clone();
if (!kSrc(cc).IsEmpty() && !kDst(cc).IsEmpty())
{
Tensor<double> _src = (*kSrc(cc));
Tensor<double> _dst = (*kDst(cc));
Mat clone_image = mat_image_.clone();
for (int i = 0; i < 854; ++i) for (int i = 0; i < 854; ++i)
{ {
if (i == 246) if (i == 246)
@ -277,9 +125,21 @@ namespace mediapipe
} }
} }
} }
out.copyTo(*image_mat); }
out.copyTo(*image);
auto output_frame = absl::make_unique<ImageFrame>(
input.Format(), mat_image_.cols, mat_image_.rows);
uchar *image_mat_ptr = image->data;
output_frame->CopyPixelData(input.Format(), mat_image_.cols, mat_image_.rows, image_mat_ptr,
ImageFrame::kDefaultAlignmentBoundary);
kOut(cc).Send(std::move(output_frame));
return absl::OkStatus(); return absl::OkStatus();
} }
REGISTER_CALCULATOR(WarpAffineCalculator); };
MEDIAPIPE_REGISTER_NODE(WarpAffineCalculator);
}
} // namespace mediapipe } // namespace mediapipe