cellpose-cpp/cyto3.h at master · ryouchinsa/cellpose-cpp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
#ifndef CYTO3CPP_H_
#define CYTO3CPP_H_

#include <list>
#include <fstream>
#include <iostream>
#include <filesystem>
#include <opencv2/core.hpp>
#include <onnxruntime_cxx_api.h>
#include "fill_voids.hpp"

class Cyto3 {
  std::unique_ptr<Ort::Session> sessionEncoder;
  Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "test"};
  Ort::SessionOptions sessionOptions[1];
  Ort::RunOptions runOptionsEncoder;
  Ort::MemoryInfo memoryInfo{Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)};
  std::vector<int64_t> inputShapeEncoder;
  std::vector<unsigned short> maskVector;
  std::vector<float> flowErrorsVector;
  bool loadingModel = false;
  bool preprocessing = false;
  bool terminating = false;

 public:
  Cyto3();
  ~Cyto3();
  bool clearLoadModelCP();
  void terminatePreprocessingCP();
  bool loadModel(const std::string& encoderPath, int threadsNumber, std::string device);
  void loadingStart();
  void loadingEnd();
  cv::Size getInputSize();
  std::tuple<cv::Mat, cv::Mat> preprocessImage(const cv::Mat& image, const cv::Size &imageSize, const std::vector<int64_t> &channels, int diameter, float cellprob_threshold, int niter, float flow_threshold, int min_size);
  cv::Mat changeFlowThreshold(float flow_threshold, int min_size);
  void preprocessingStart();
  void preprocessingEnd();
};

void saveOutputMask(cv::Mat mask, cv::Size imageSize, std::string path);

#endif