↰ Return to documentation for file (morpheus/_lib/src/stages/add_scores.cpp
)
#include "morpheus/stages/add_scores.hpp"
#include "morpheus/objects/tensor.hpp"
#include "morpheus/objects/tensor_object.hpp"// for TensorIndex, TensorObject
#include <glog/logging.h>
#include <mrc/node/sink_properties.hpp>// for SinkProperties<>::sink_type_t
#include <mrc/node/source_properties.hpp>// for SourceProperties<>::source_type_t
#include <mrc/segment/object.hpp>// for Object
#include <cstddef>// for size_t
#include <exception>
#include <map>
#include <memory>
#include <ostream>// for logging
#include <type_traits>// for declval
#include <utility>// for move
#include <vector>
// IWYU thinks we need __alloc_traits<>::value_type for vector assignments
// IWYU pragma: no_include <ext/alloc_traits.h>
namespace morpheus {
// Component public implementations
// ************ AddScoresStage **************************** //
AddScoresStage::AddScoresStage(std::size_t num_class_labels, std::map<std::size_t, std::string> idx2label) :
PythonNode(base_t::op_factory_from_sub_fn(build_operator())),
m_num_class_labels(num_class_labels),
m_idx2label(std::move(idx2label))
{
CHECK(m_idx2label.size() <= m_num_class_labels) << "idx2label should represent a subset of the class_labels";
}
AddScoresStage::subscribe_fn_t AddScoresStage::build_operator()
{
return [this](rxcpp::observable<sink_type_t> input, rxcpp::subscriber<source_type_t> output) {
return input.subscribe(rxcpp::make_observer<sink_type_t>(
[this, &output](sink_type_t x) {
const auto& probs = x->get_probs();
const auto& shape = probs.get_shape();
CHECK(shape.size() == 2 && shape[1] == m_num_class_labels)
<< "Label count does not match output of model. Label count: " << m_num_class_labels
<< ", Model output: " << shape[1];
const std::size_t num_rows = shape[0];
const std::size_t num_columns = shape[1];
std::vector<std::string> columns(m_idx2label.size());
std::vector<TensorObject> tensors(m_idx2label.size());
std::size_t i = 0;
for (const auto& [column_num, column_name] : m_idx2label)
{
columns[i] = column_name;
tensors[i] = probs.slice(std::vector<TensorIndex>{0, static_cast<TensorIndex>(column_num)},
std::vector<TensorIndex>{static_cast<TensorIndex>(num_rows),
static_cast<TensorIndex>(column_num + 1)});
++i;
}
x->set_meta(columns, tensors);
output.on_next(x);
},
[&](std::exception_ptr error_ptr) { output.on_error(error_ptr); },
[&]() { output.on_completed(); }));
};
}
// ************ AddScoresStageInterfaceProxy ************* //
std::shared_ptr<mrc::segment::Object<AddScoresStage>> AddScoresStageInterfaceProxy::init(
mrc::segment::Builder& builder,
const std::string& name,
std::size_t num_class_labels,
std::map<std::size_t, std::string> idx2label)
{
return builder.construct_object<AddScoresStage>(name, num_class_labels, std::move(idx2label));
}
} // namespace morpheus