cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
2044e45b1a
mediapipe/mediapipe2/util/tracking/tone_models.h
AK391 2044e45b1a examples
2021-06-10 23:01:19 +00:00

629 lines
21 KiB
C++
Raw Blame History

// 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.
//
// Implements various tone-models describing tonal change of color intensities
// across frame pairs.
#ifndef MEDIAPIPE_UTIL_TRACKING_TONE_MODELS_H_
#define MEDIAPIPE_UTIL_TRACKING_TONE_MODELS_H_
#include <algorithm>
#include <cmath>
#include <string>
#include <vector>
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
#include "mediapipe/framework/port/vector.h"
#include "mediapipe/util/tracking/tone_models.pb.h"
namespace mediapipe {
// Abstract Adapter for tone models.
template <class Model>
class ToneModelAdapter {
public:
// Initialized Model from pointer to arguments. If identitiy_parametrization
// is set, args = {0,0, ... 0} correspond to identity model.
template <class T>
static Model FromPointer(const T* args, bool identity_parametrization);
// Outputs parameters to pointer args.
template <class T>
static void ToPointer(const GainBiasModel& model, T* args);
// Transforms color (RGB space) according to model.
static Vector3_f TransformColor(const Model& model, const Vector3_f& color);
// Inverts color mode, if not invertible success is set to zero.
static Model InvertChecked(const Model& model, bool* success);
static Model Compose(const Model& lhs, const Model& rhs);
static int NumParameters();
static float GetParameter(const Model& model, int idx);
static float Determinant(const Model& model);
};
// LogDomain LUT:
// Converts from intensity domain to log-domain via Map, inverse mapping via
// UnMap.
// Use as Singleton only, via inline LogDomainLUT()
class LogDomainLUTImpl {
public:
LogDomainLUTImpl(const LogDomainLUTImpl&) = delete;
LogDomainLUTImpl& operator=(const LogDomainLUTImpl&) = delete;
// Maps intensity between [0, 255] to log-domain by applying log(1 + x).
// Result will be within [0, log(256)]. Truncation occurs for non-integer
// inputs. Discritization error is at most 0.5 intensity levels.
inline float Map(float value) const {
return log_lut_[std::max(0, std::min<int>(255, value + 0.5f))];
}
inline Vector3_f MapVec(Vector3_f vec) const {
return Vector3_f(Map(vec.x()), Map(vec.y()), Map(vec.z()));
}
// Unmaps intensity from log-domain in [0, log(256)] = [0, 2.4] to regular
// domain, via exp(x) - 1.0. Discritization error of at most 0.01 intensity
// levels might occur.
inline float UnMap(float value) const {
const int bin = value * inv_max_log_value_ * (kExpBins - 2);
return exp_lut_[std::max(0, std::min(bin, kExpBins - 1))];
}
inline Vector3_f UnMapVec(Vector3_f vec) const {
return Vector3_f(UnMap(vec.x()), UnMap(vec.y()), UnMap(vec.z()));
}
float MaxLogDomainValue() const { return max_log_value_; } // ~ log(256).
private:
friend const LogDomainLUTImpl& LogDomainLUT();
LogDomainLUTImpl();
std::vector<float> log_lut_;
std::vector<float> exp_lut_;
float max_log_value_;
float inv_max_log_value_;
enum {
kExpBins = 2560
}; // Allocated such that exp(2.4) -
// exp(2.4 * (1.0 - 1.0 / kExpBins)) < 0.01.
};
inline const LogDomainLUTImpl& LogDomainLUT() {
static auto p = new LogDomainLUTImpl();
return *p;
}
inline Vector3_i RoundAndClampColor(const Vector3_f& vec) {
return Vector3_i(
std::max(0, std::min(255, static_cast<int>(vec.x() + 0.5f))),
std::max(0, std::min(255, static_cast<int>(vec.y() + 0.5f))),
std::max(0, std::min(255, static_cast<int>(vec.z() + 0.5f))));
}
template <>
class ToneModelAdapter<GainBiasModel> {
public:
template <class T>
static GainBiasModel FromPointer(const T* args,
bool identity_parametrization);
template <class T>
static void ToPointer(const GainBiasModel& model, T* args);
inline static Vector3_f TransformColor(const GainBiasModel& model,
const Vector3_f& color);
inline static GainBiasModel InvertChecked(const GainBiasModel& model,
bool* success);
inline static GainBiasModel Compose(const GainBiasModel& lhs,
const GainBiasModel& rhs);
enum { kNumParameters = 6 };
static int NumParameters() { return kNumParameters; }
inline static float GetParameter(const GainBiasModel& model, int idx);
static float Determinant(const GainBiasModel& model);
static GainBiasModel AddIdentity(const GainBiasModel& model);
static GainBiasModel ScaleParameters(const GainBiasModel& model, float scale);
static std::string ToString(const GainBiasModel& model);
};
template <>
class ToneModelAdapter<AffineToneModel> {
public:
template <class T>
static AffineToneModel FromPointer(const T* args,
bool identity_parametrization);
template <class T>
static void ToPointer(const AffineToneModel& model, T* args) {
return ToPointerPad(model, false, args);
}
// If pad_square is set the 12 DOF 3x4 model is embedded in a 4x4 model
// (last row = [ 0 0 0 1]) and output to args (row major ordering),
// otherwise the original 12 DOF 3x4 model is written to args (row major).
template <class T>
static void ToPointerPad(const AffineToneModel& model, bool pad_square,
T* args);
inline static Vector3_f TransformColor(const AffineToneModel& model,
const Vector3_f& color);
inline static AffineToneModel InvertChecked(const AffineToneModel& model,
bool* success);
inline static AffineToneModel Compose(const AffineToneModel& lhs,
const AffineToneModel& rhs);
enum { kNumParameters = 12 };
static int NumParameters() { return kNumParameters; }
inline static float GetParameter(const AffineToneModel& model, int idx);
// Used during stabilization: Adds identity model (model + I)
// and scales each parameter by scale.
static AffineToneModel AddIdentity(const AffineToneModel& model);
static AffineToneModel ScaleParameters(const AffineToneModel& model,
float scale);
};
typedef ToneModelAdapter<GainBiasModel> GainBiasModelAdapter;
typedef ToneModelAdapter<AffineToneModel> AffineToneModelAdapter;
// Mixture Models.
template <class BaseModel, class MixtureModel>
struct MixtureToneTraits {
typedef BaseModel BaseModelType;
typedef MixtureModel ModelType;
};
typedef MixtureToneTraits<GainBiasModel, MixtureGainBiasModel>
GainBiasModelTraits;
typedef MixtureToneTraits<AffineToneModel, MixtureAffineToneModel>
AffineToneModelTraits;
template <class Traits>
class MixtureToneAdapter {
typedef typename Traits::ModelType MixtureModel;
typedef typename Traits::BaseModelType BaseModel;
typedef ToneModelAdapter<BaseModel> BaseModelAdapter;
enum { kBaseNumParameters = BaseModelAdapter::kNumParameters };
public:
template <class T>
static MixtureModel FromPointer(const T* args, bool identity_parametrization,
int skip, int num_models);
template <class T>
static void ToPointer(const MixtureModel& model, T* args);
static int NumParameters(const MixtureModel& model) {
return model.model_size() * kBaseNumParameters;
}
inline static float GetParameter(const MixtureModel& model, int model_id,
int param_id);
inline static MixtureModel IdentityModel(int num_mixtures);
inline static BaseModel ToBaseModel(const MixtureModel& mixture_model,
const float* weights);
inline static Vector3_f TransformColor(const MixtureModel& model,
const float* weights,
const Vector3_f& pt);
// Mixtures are not invertible via closed form, but invertible for a specific
// set of weights by computing the underlying BaseModel via ToBaseModel.
// Resulting BaseModel is applied to pt, if BaseModel is not invertible
// success is set to false.
inline static Vector3_f SolveForColorChecked(const MixtureModel& model,
const float* weights,
const Vector3_f& pt,
bool* success);
};
typedef MixtureToneAdapter<GainBiasModelTraits> MixtureGainBiasModelAdapter;
typedef MixtureToneAdapter<AffineToneModelTraits> MixtureAffineToneModelAdapter;
template <class T>
GainBiasModel ToneModelAdapter<GainBiasModel>::FromPointer(const T* args,
bool identity) {
DCHECK(args);
GainBiasModel model;
const float id_shift = identity ? 1.0f : 0.0f;
model.set_gain_c1(args[0] + id_shift);
model.set_bias_c1(args[1]);
model.set_gain_c2(args[2] + id_shift);
model.set_bias_c2(args[3]);
model.set_gain_c3(args[4] + id_shift);
model.set_bias_c3(args[5]);
return model;
}
template <class T>
void ToneModelAdapter<GainBiasModel>::ToPointer(const GainBiasModel& model,
T* args) {
args[0] = model.gain_c1();
args[1] = model.bias_c1();
args[2] = model.gain_c2();
args[3] = model.bias_c2();
args[4] = model.gain_c3();
args[5] = model.bias_c3();
}
inline Vector3_f ToneModelAdapter<GainBiasModel>::TransformColor(
const GainBiasModel& model, const Vector3_f& color) {
return Vector3_f(model.gain_c1() * color.x() + model.bias_c1(),
model.gain_c2() * color.y() + model.bias_c2(),
model.gain_c3() * color.z() + model.bias_c3());
}
inline float ToneModelAdapter<GainBiasModel>::Determinant(
const GainBiasModel& model) {
return model.gain_c1() * model.gain_c2() * model.gain_c3();
}
inline GainBiasModel ToneModelAdapter<GainBiasModel>::InvertChecked(
const GainBiasModel& model, bool* success) {
// (g_1 0 0 b_1 == (1/g_1 0 0 -b_1/g_1
// 0 g_2 0 b_2 0 1/g_2 0 -b_2/g_2
// 0 0 g_3 b_3 0 0 1/g_3 -b_3/g_3
// 0 0 0 1)^(-1) 0 0 0 1
// Compute determinant.
const float det = GainBiasModelAdapter::Determinant(model);
if (fabs(det) < 1e-10f) {
*success = false;
LOG(ERROR) << "Model not invertible.";
return GainBiasModel();
}
*success = true;
GainBiasModel result;
const float inv_gain_c1 = 1.0f / model.gain_c1();
const float inv_gain_c2 = 1.0f / model.gain_c2();
const float inv_gain_c3 = 1.0f / model.gain_c3();
result.set_gain_c1(inv_gain_c1);
result.set_bias_c1(inv_gain_c1 * -model.bias_c1());
result.set_gain_c2(inv_gain_c2);
result.set_bias_c2(inv_gain_c2 * -model.bias_c2());
result.set_gain_c3(inv_gain_c3);
result.set_bias_c3(inv_gain_c3 * -model.bias_c3());
return result;
}
inline GainBiasModel ToneModelAdapter<GainBiasModel>::Compose(
const GainBiasModel& lhs, const GainBiasModel& rhs) {
GainBiasModel result;
result.set_gain_c1(lhs.gain_c1() * rhs.gain_c1());
result.set_bias_c1(lhs.gain_c1() * rhs.bias_c1() + lhs.bias_c1());
result.set_gain_c2(lhs.gain_c2() * rhs.gain_c2());
result.set_bias_c2(lhs.gain_c2() * rhs.bias_c2() + lhs.bias_c2());
result.set_gain_c3(lhs.gain_c3() * rhs.gain_c3());
result.set_bias_c3(lhs.gain_c3() * rhs.bias_c3() + lhs.bias_c3());
return result;
}
inline float ToneModelAdapter<GainBiasModel>::GetParameter(
const GainBiasModel& model, int idx) {
switch (idx) {
case 0:
return model.gain_c1();
case 1:
return model.bias_c1();
case 2:
return model.gain_c2();
case 3:
return model.bias_c2();
case 4:
return model.gain_c3();
case 5:
return model.bias_c3();
default:
LOG(FATAL) << "Unknown parameter requested.";
}
return 0.0f;
}
template <class T>
AffineToneModel ToneModelAdapter<AffineToneModel>::FromPointer(const T* args,
bool identity) {
DCHECK(args);
AffineToneModel model;
const float id_shift = identity ? 1.0f : 0.0f;
model.set_g_00(args[0] + id_shift);
model.set_g_01(args[1]);
model.set_g_02(args[2]);
model.set_g_03(args[3]);
model.set_g_10(args[4]);
model.set_g_11(args[5] + id_shift);
model.set_g_12(args[6]);
model.set_g_13(args[7]);
model.set_g_20(args[8]);
model.set_g_21(args[9]);
model.set_g_22(args[10] + id_shift);
model.set_g_23(args[11]);
return model;
}
template <class T>
void ToneModelAdapter<AffineToneModel>::ToPointerPad(
const AffineToneModel& model, bool pad_square, T* args) {
DCHECK(args);
args[0] = model.g_00();
args[1] = model.g_01();
args[2] = model.g_02();
args[3] = model.g_03();
args[4] = model.g_10();
args[5] = model.g_11();
args[6] = model.g_12();
args[7] = model.g_13();
args[8] = model.g_20();
args[9] = model.g_21();
args[10] = model.g_22();
args[11] = model.g_23();
if (pad_square) {
args[12] = 0;
args[13] = 0;
args[14] = 0;
args[15] = 1;
}
}
inline Vector3_f ToneModelAdapter<AffineToneModel>::TransformColor(
const AffineToneModel& model, const Vector3_f& color) {
return Vector3_f(model.g_00() * color.x() + model.g_01() * color.y() +
model.g_02() * color.z() + model.g_03(),
model.g_10() * color.x() + model.g_11() * color.y() +
model.g_12() * color.z() + model.g_13(),
model.g_20() * color.x() + model.g_21() * color.y() +
model.g_22() * color.z() + model.g_23());
}
inline AffineToneModel ToneModelAdapter<AffineToneModel>::InvertChecked(
const AffineToneModel& model, bool* success) {
double data[16];
double inv_data[16];
ToPointerPad<double>(model, true, data);
cv::Mat model_mat(4, 4, CV_64F, data);
cv::Mat inv_model_mat(4, 4, CV_64F, inv_data);
if (cv::invert(model_mat, inv_model_mat) < 1e-10) {
LOG(ERROR) << "AffineToneModel not invertible, det is zero.";
*success = false;
return AffineToneModel();
}
*success = true;
return FromPointer<double>(inv_data, false);
}
inline AffineToneModel ToneModelAdapter<AffineToneModel>::Compose(
const AffineToneModel& lhs, const AffineToneModel& rhs) {
AffineToneModel result;
double lhs_data[16];
double rhs_data[16];
double result_data[16];
ToPointerPad<double>(lhs, true, lhs_data);
ToPointerPad<double>(rhs, true, rhs_data);
cv::Mat lhs_mat(4, 4, CV_64F, lhs_data);
cv::Mat rhs_mat(4, 4, CV_64F, rhs_data);
cv::Mat result_mat(4, 4, CV_64F, result_data);
cv::gemm(lhs_mat, rhs_mat, 1.0, cv::Mat(), 0, result_mat);
return FromPointer<double>(result_data, false);
}
inline float ToneModelAdapter<AffineToneModel>::GetParameter(
const AffineToneModel& model, int idx) {
switch (idx) {
case 0:
return model.g_00();
case 1:
return model.g_01();
case 2:
return model.g_02();
case 3:
return model.g_03();
case 4:
return model.g_10();
case 5:
return model.g_11();
case 6:
return model.g_12();
case 7:
return model.g_13();
case 8:
return model.g_20();
case 9:
return model.g_21();
case 10:
return model.g_22();
case 11:
return model.g_23();
default:
LOG(FATAL) << "Unknown parameter requested.";
}
return 0.0f;
}
template <class Traits>
template <class T>
typename Traits::ModelType MixtureToneAdapter<Traits>::FromPointer(
const T* args, bool identity_parametrization, int skip, int num_models) {
MixtureModel model;
const T* arg_ptr = args;
for (int i = 0; i < num_models; ++i, arg_ptr += kBaseNumParameters + skip) {
BaseModel base =
BaseModelAdapter::FromPointer(arg_ptr, identity_parametrization);
model.add_model()->CopyFrom(base);
}
return model;
}
template <class Traits>
template <class T>
void MixtureToneAdapter<Traits>::ToPointer(const MixtureModel& model, T* args) {
const T* arg_ptr = args;
const int num_models = model.model_size();
for (int i = 0; i < num_models; ++i, arg_ptr += kBaseNumParameters) {
BaseModel base = BaseModelAdapter::ToPointer(model.model(i), arg_ptr);
}
}
template <class Traits>
inline float MixtureToneAdapter<Traits>::GetParameter(const MixtureModel& model,
int model_id,
int param_id) {
return BaseModelAdapter::GetParameter(model.model(model_id), param_id);
}
template <class Traits>
inline typename Traits::ModelType MixtureToneAdapter<Traits>::IdentityModel(
int num_mixtures) {
MixtureModel model;
for (int i = 0; i < num_mixtures; ++i) {
model.add_model();
}
return model;
}
template <class Traits>
inline typename MixtureToneAdapter<Traits>::BaseModel
MixtureToneAdapter<Traits>::ToBaseModel(const MixtureModel& mixture_model,
const float* weights) {
const int num_models = mixture_model.model_size();
float params[kBaseNumParameters];
memset(params, 0, sizeof(params[0]) * kBaseNumParameters);
// Weighted combination of mixture models.
for (int m = 0; m < num_models; ++m) {
for (int k = 0; k < kBaseNumParameters; ++k) {
params[k] += BaseModelAdapter::GetParameter(mixture_model.model(m), k) *
weights[m];
}
}
return BaseModelAdapter::FromPointer(params, false);
}
template <class Traits>
inline Vector3_f MixtureToneAdapter<Traits>::TransformColor(
const MixtureModel& model, const float* weights, const Vector3_f& pt) {
const int num_models = model.model_size();
Vector3_f result(0, 0, 0);
for (int i = 0; i < num_models; ++i) {
result += BaseModelAdapter::TransformColor(model.model(i), pt) * weights[i];
}
return result;
}
template <class Traits>
inline Vector3_f MixtureToneAdapter<Traits>::SolveForColorChecked(
const MixtureModel& model, const float* weights, const Vector3_f& pt,
bool* success) {
BaseModel base_model = ToBaseModel(model, weights);
BaseModel inv_base_model =
BaseModelAdapter::InvertChecked(base_model, success);
return BaseModelAdapter::TransformColor(inv_base_model, pt);
}
template <class Model, class Adapter>
class ToneModelMethods {
public:
// Maps image input to output by applying model to each pixel's intensity.
// Set log_domain to true, if model has been estimated in the log-domain.
// Set normalized_model to true, if model is normalized (expected input
// intensity range in [0, 1]).
// Number of output channels (N) should be <= number of input channels (M),
// in which case the only first N input channels are copied to the first N
// output channels.
static void MapImage(const Model& model, bool log_domain,
bool normalized_model, const cv::Mat& input,
cv::Mat* output);
// Fast mapping version of above function via LUT if color model is
// independent across specified C channels (that is model is a diagonal
// matrix).
// Otherwise incorrect results will be obtained.
template <int C>
static void MapImageIndependent(const Model& model, bool log_domain,
bool normalized_model, const cv::Mat& input,
cv::Mat* output);
};
typedef ToneModelMethods<GainBiasModel, GainBiasModelAdapter>
GainBiasModelMethods;
typedef ToneModelMethods<AffineToneModel, AffineToneModelAdapter>
AffineToneModelMethods;
template <class Model, class Adapter>
template <int C>
void ToneModelMethods<Model, Adapter>::MapImageIndependent(
const Model& model, bool log_domain, bool normalized_model,
const cv::Mat& input, cv::Mat* output) {
CHECK(output != nullptr);
CHECK_EQ(input.channels(), C);
CHECK_EQ(output->channels(), C);
// Input LUT which will be mapped to the output LUT by the tone change model.
// Needs 3 channels to represent input RGB colors, but since they are assumed
// independent, we can assign the same value to each channel, which will be
// transformed by the respective channel's transform.
cv::Mat lut_input(1, 256, CV_8UC3);
uint8* lut_ptr = lut_input.ptr<uint8>(0);
for (int k = 0; k < 256; ++k, lut_ptr += 3) {
for (int c = 0; c < 3; ++c) {
lut_ptr[c] = k;
}
}
// Output LUT. Only needs C channels (as many in the input/output).
cv::Mat lut(1, 256, CV_8UC(C));
MapImage(model, log_domain, normalized_model, lut_input, &lut);
cv::LUT(input, lut, *output);
}
// TODO: Implement for mixtures.
// typedef ToneModelMethods<MixtureGainBiasModel,
// MixtureGainBiasModelAdapter>
// MixtureGainBiasModelMethods;
// typedef ToneModelMethods<MixtureAffineToneModel,
// MixtureAffineToneModelAdapter>
// MixtureAffineToneModelMethods;
} // namespace mediapipe
#endif // MEDIAPIPE_UTIL_TRACKING_TONE_MODELS_H_
Reference in New Issue View Git Blame Copy Permalink