| 158 | |
| 159 | template <class Model> |
| 160 | void Segmentor<Model>::Predict(cv::Mat image, std::string which_class){ |
| 161 | cv::Mat srcImg = image.clone(); |
| 162 | int which_class_index = -1; |
| 163 | for(int i = 0; i<name_list.size(); i++){ |
| 164 | if(name_list[i] == which_class){ |
| 165 | which_class_index = i; |
| 166 | break; |
| 167 | } |
| 168 | } |
| 169 | if(which_class_index==-1) throw which_class + "not in the name list"; |
| 170 | int image_width = image.cols; |
| 171 | int image_height = image.rows; |
| 172 | cv::resize(image,image,cv::Size(width,height)); |
| 173 | torch::Tensor tensor_image = torch::from_blob(image.data, { 1, height, width,3 }, torch::kByte); |
| 174 | tensor_image = tensor_image.to(device); |
| 175 | tensor_image = tensor_image.permute({ 0,3,1,2 }); |
| 176 | tensor_image = tensor_image.to(torch::kFloat); |
| 177 | tensor_image = tensor_image.div(255.0); |
| 178 | |
| 179 | at::Tensor output = model->forward({ tensor_image }); |
| 180 | output = torch::softmax(output, 1).mul(255.0).toType(torch::kByte); |
| 181 | |
| 182 | image = cv::Mat::ones(cv::Size(width, height), CV_8UC1); |
| 183 | |
| 184 | at::Tensor re = output[0][which_class_index].to(at::kCPU).detach(); |
| 185 | memcpy(image.data, re.data_ptr(), width * height * sizeof(unsigned char)); |
| 186 | cv::resize(image, image, cv::Size(image_width, image_height)); |
| 187 | |
| 188 | // draw the prediction |
| 189 | cv::imwrite("prediction.jpg", image); |
| 190 | cv::imshow("prediction", image); |
| 191 | cv::imshow("srcImage", srcImg); |
| 192 | cv::waitKey(0); |
| 193 | cv::destroyAllWindows(); |
| 194 | return; |
| 195 | } |
| 196 | |
| 197 | |
| 198 | #endif // SEGMENTOR_H |