MLSegmentation.h

Go to the documentation of this file.

#ifndef _SEIMPLEMENTATION_SEGMENTATION_MLSEGMENTATION_H_
#define _SEIMPLEMENTATION_SEGMENTATION_MLSEGMENTATION_H_

#include "SEFramework/Source/SourceFactory.h"
#include "SEFramework/Pipeline/Segmentation.h"

namespace SourceXtractor {


class MLSegmentation : public Segmentation::Labelling {
public:

  virtual ~MLSegmentation() = default;

  MLSegmentation(std::shared_ptr<SourceFactory> source_factory, std::string model_path, double ml_threshold)
      : m_source_factory(source_factory), m_model_path(model_path), m_ml_threshold(ml_threshold) {
    assert(source_factory != nullptr);
  }

  void labelImage(Segmentation::LabellingListener& listener, std::shared_ptr<const DetectionImageFrame> frame) override;

private:
  struct Tile {
    PixelCoordinate offset;
    int width, height;
  };

  std::vector<MLSegmentation::Tile> getTiles(const DetectionImage& image) const;


  std::shared_ptr<SourceFactory> m_source_factory;

  std::string m_model_path;
  double m_ml_threshold;

};

}

#endif /* _SEIMPLEMENTATION_SEGMENTATION_MLSEGMENTATION_H_ */