| 146 | } |
| 147 | |
| 148 | int Classifier::Predict(cv::Mat& image){ |
| 149 | cv::resize(image, image, cv::Size(448, 448)); |
| 150 | torch::Tensor img_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }); |
| 151 | img_tensor = img_tensor.to(device).unsqueeze(0).to(torch::kF32).div(255.0); |
| 152 | auto prediction = vgg->forward(img_tensor); |
| 153 | prediction = torch::softmax(prediction,1); |
| 154 | auto class_id = prediction.argmax(1); |
| 155 | std::cout<<prediction<<class_id; |
| 156 | int ans = int(class_id.item().toInt()); |
| 157 | float prob = prediction[0][ans].item().toFloat(); |
| 158 | return ans; |
| 159 | } |
| 160 | |
| 161 | void Classifier::LoadWeight(std::string weight){ |
| 162 | torch::load(vgg,weight); |