Add BGR -> RGB color conversion to ColorConvertCalculator.

PiperOrigin-RevId: 496497002
This commit is contained in:
MediaPipe Team 2022-12-19 14:48:47 -08:00 committed by Copybara-Service
parent ef3fa67bf4
commit ea0bebc226

View File

@ -38,6 +38,7 @@ void SetColorChannel(int channel, uint8 value, cv::Mat* mat) {
constexpr char kRgbaInTag[] = "RGBA_IN";
constexpr char kRgbInTag[] = "RGB_IN";
constexpr char kBgrInTag[] = "BGR_IN";
constexpr char kBgraInTag[] = "BGRA_IN";
constexpr char kGrayInTag[] = "GRAY_IN";
constexpr char kRgbaOutTag[] = "RGBA_OUT";
@ -57,6 +58,7 @@ constexpr char kGrayOutTag[] = "GRAY_OUT";
// RGB -> RGBA
// RGBA -> BGRA
// BGRA -> RGBA
// BGR -> RGB
//
// This calculator only supports a single input stream and output stream at a
// time. If more than one input stream or output stream is present, the
@ -69,6 +71,7 @@ constexpr char kGrayOutTag[] = "GRAY_OUT";
// RGB_IN: The input video stream (ImageFrame, SRGB).
// BGRA_IN: The input video stream (ImageFrame, SBGRA).
// GRAY_IN: The input video stream (ImageFrame, GRAY8).
// BGR_IN: The input video stream (ImageFrame, SBGR).
//
// Output streams:
// RGBA_OUT: The output video stream (ImageFrame, SRGBA).
@ -122,6 +125,10 @@ absl::Status ColorConvertCalculator::GetContract(CalculatorContract* cc) {
cc->Inputs().Tag(kBgraInTag).Set<ImageFrame>();
}
if (cc->Inputs().HasTag(kBgrInTag)) {
cc->Inputs().Tag(kBgrInTag).Set<ImageFrame>();
}
if (cc->Outputs().HasTag(kRgbOutTag)) {
cc->Outputs().Tag(kRgbOutTag).Set<ImageFrame>();
}
@ -194,6 +201,11 @@ absl::Status ColorConvertCalculator::Process(CalculatorContext* cc) {
return ConvertAndOutput(kRgbaInTag, kBgraOutTag, ImageFormat::SBGRA,
cv::COLOR_RGBA2BGRA, cc);
}
// BGR -> RGB
if (cc->Inputs().HasTag(kBgrInTag) && cc->Outputs().HasTag(kRgbOutTag)) {
return ConvertAndOutput(kBgrInTag, kRgbOutTag, ImageFormat::SRGB,
cv::COLOR_BGR2RGB, cc);
}
return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "Unsupported image format conversion.";