-
Notifications
You must be signed in to change notification settings - Fork 1
/
brainUnet.hpp
37 lines (29 loc) · 942 Bytes
/
brainUnet.hpp
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
#pragma once
#include<iostream>
#include<vector>
#include<torch/script.h>
#include<opencv2/opencv.hpp>
#include<time.h>
#include<map>
using std::cout;
using std::endl;
//Meyers Singleton pattern
class BrainUnetModel{
public:
static BrainUnetModel& getInstance();
cv::Mat toOpencvImage(torch::Tensor& pred_img);
void applyMaskRoi(cv::Mat& original_img,cv::Mat& predicted_mask) const;
void showCountersOfTumor(std::tuple<cv::Mat, cv::Mat>rgb_and_mask_image) const;
cv::Mat forwardModel();
void setModelPath(const std::string& model_path);
void setDataPath(const std::string& data_path);
private:
~BrainUnetModel()= default;
BrainUnetModel() = default;
BrainUnetModel(const BrainUnetModel&) = delete;
BrainUnetModel& operator = (const BrainUnetModel&) = delete;
cv::Mat m_input_img;
cv::Mat m_rgb_img;
cv::Mat m_predicted_mask;
torch::jit::Module m_module;
};