Class SegmentationParser

Inheritance Relationships

Base Types

Class Documentation

class SegmentationParser : public dai::DeviceNodeCRTP<DeviceNode, SegmentationParser, SegmentationParserProperties>, public dai::HostRunnable

SegmentationParser node. Parses raw segmentation output from segmentation neural networks into a dai::SegmentationMask datatype. The parser supports two output model types:

  1. Single-channel output where the model argmaxes the class probabilities internally and outputs a single channel mask with class indices.

  2. Multi-channel output where each channel corresponds to the probability map for a specific class. The parser will perform argmax across channels to generate the final mask. The parser can be configured to treat the first class (index 0) as the background class, which will be ignored in the final segmentation mask.

Warning

Only OAK4 supports running SegmentationParser on device. On other platforms, the node will automatically switch to host execution.

Public Types

using Model = std::variant<NNModelDescription, NNArchive, std::string>

Public Functions

~SegmentationParser() override
SegmentationParser() = default
inline SegmentationParser(std::unique_ptr<Properties> props)
std::shared_ptr<SegmentationParser> build(Node::Output &nnInput, const Model &model)

Build SegmentationParser node.

Parameters:

  • nnInput – Output to link

  • nnArchive – Neural network archive

std::shared_ptr<SegmentationParser> build(Node::Output &nnInput, const dai::nn_archive::v1::Head &head)

Build SegmentationParser node with the specific head from NNArchive. Useful when model outputs multiple segmentation heads.

Parameters:

  • nnInput – Output to link

  • head – Specific head from NNArchive to use for this parser

void setLabels(const std::vector<std::string> &labels)

Sets the class labels associated with the segmentation mask. The label at index $i$ in the labels vector corresponds to the value $i$ in the segmentation mask data array.

Parameters:

labels – Vector of class labels

std::vector<std::string> getLabels() const

Returns the class labels associated with the segmentation mask.

void setBackgroundClass(bool backgroundClass)

Sets whether the first class (index 0) is considered the background class. If true, the pixels classified as index 0 will be treated as background.

Note

Only applicable if the number of classes is greater than 1 and the output classes are not in a single layer (eg. classesInOneLayer = false).

Parameters:

backgroundClass – Boolean indicating if the first class is the background class

bool getBackgroundClass() const

Gets whether the first class (index 0) is considered the background class.

void setRunOnHost(bool runOnHost)

Specify whether to run on host or device By default, the node will run on device.

virtual bool runOnHost() const override

Check if the node is set to run on host

virtual void run() override
virtual void buildInternal() override

Function called from within the create function to build the node. This function is useful for initialization, setting up inputs and outputs = stuff that cannot be perform in the constuctor.

inline DeviceNodeCRTP()
inline DeviceNodeCRTP(const std::shared_ptr<Device> &device)
inline DeviceNodeCRTP(std::unique_ptr<Properties> props)
inline DeviceNodeCRTP(std::unique_ptr<Properties> props, bool confMode)
inline DeviceNodeCRTP(const std::shared_ptr<Device> &device, std::unique_ptr<Properties> props, bool confMode)

Public Members

std::shared_ptr<SegmentationParserConfig> initialConfig = std::make_shared<SegmentationParserConfig>()

Initial config to use when parsing segmentation masks.

Input input = {*this, {"input", DEFAULT_GROUP, true, 5, {{{DatatypeEnum::NNData, true}}}, true}}

Input NN results with segmentation data to parser

Input inputConfig = {*this, {"inputConfig", DEFAULT_GROUP, false, 4, {{{DatatypeEnum::SegmentationParserConfig, false}}}, DEFAULT_WAIT_FOR_MESSAGE}}

Input SegmentationParserConfig message with ability to modify parameters in runtime.

Output out = {*this, {"out", DEFAULT_GROUP, {{{DatatypeEnum::SegmentationMask, false}}}}}

Outputs segmentation mask

Public Static Attributes

static constexpr const char *NAME = "SegmentationParser"

Protected Functions

virtual Properties &getProperties() override