| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| #include "result.h" |
|
|
| #include <fstream> |
| #include <string> |
|
|
| #include "src/utils/utility.h" |
| #include "third_party/nlohmann/json.hpp" |
|
|
| using json = nlohmann::json; |
| void DocTrResult::SaveToImg(const std::string &save_path) { |
| absl::StatusOr<std::string> full_path; |
| if (predictor_result_.input_path.empty()) { |
| auto now = std::chrono::system_clock::now(); |
| auto now_time = std::chrono::system_clock::to_time_t(now); |
| std::stringstream ss; |
| ss << "output_" << std::put_time(std::localtime(&now_time), "%Y%m%d_%H%M%S") |
| << ".jpg"; |
| std::string timestamp_filename = ss.str(); |
| INFOW("Input path is empty, will use %s instead!", |
| timestamp_filename.c_str()); |
| predictor_result_.input_path = timestamp_filename; |
| full_path = |
| Utility::SmartCreateDirectoryForImage(save_path, timestamp_filename); |
| } else { |
| full_path = Utility::SmartCreateDirectoryForImage( |
| save_path, predictor_result_.input_path); |
| } |
| bool success = cv::imwrite(full_path.value(), predictor_result_.doctr_img); |
| if (!success) { |
| INFOE("Failed to write the image %s", full_path.value().c_str()); |
| exit(-1); |
| } |
| } |
|
|
| void DocTrResult::Print() const { |
| std::cout << "{\n \"res\": {" << std::endl; |
|
|
| std::cout << " \"input_path\": {" << predictor_result_.input_path << "}," |
| << std::endl; |
| std::cout << " \"doctr_img\": {" |
| << "..." |
| << "}" << std::endl; |
| std::cout << "}" << std::endl; |
| } |
|
|
| void DocTrResult::SaveToJson(const std::string &save_path) const { |
| nlohmann::ordered_json j; |
|
|
| j["input_path"] = predictor_result_.input_path; |
| j["page_index"] = nullptr; |
|
|
| nlohmann::json mat_array = nlohmann::json::array(); |
|
|
| for (int i = 0; i < predictor_result_.doctr_img.rows; ++i) { |
| nlohmann::json row = nlohmann::json::array(); |
| for (int j = 0; j < predictor_result_.doctr_img.cols; ++j) { |
| cv::Vec3b color = predictor_result_.doctr_img.at<cv::Vec3b>(i, j); |
| row.push_back({color[0], color[1], color[2]}); |
| } |
| mat_array.push_back(row); |
| } |
| j["doctr_img"] = mat_array; |
|
|
| auto full_path = Utility::SmartCreateDirectoryForJson( |
| save_path, predictor_result_.input_path); |
|
|
| if (!full_path.ok()) { |
| INFOE(full_path.status().ToString().c_str()); |
| exit(-1); |
| } |
| std::ofstream file(full_path.value()); |
| if (file.is_open()) { |
| file << j.dump(4); |
| file.close(); |
| } else { |
| INFOE("Could not open file for writing: %s", save_path.c_str()); |
| exit(-1); |
| } |
| } |
|
|
| int DocTrResult::getAdaptiveFontScale(const std::string &text, int imgWidth, |
| int maxWidth, int minFont, int maxFont, |
| int thickness, int &outBaseline, |
| int &outFontFace) { |
| int fontFace = cv::FONT_HERSHEY_SIMPLEX; |
| double fontScale = 1.0; |
| int baseline = 0; |
| int bestFontSize = minFont; |
|
|
| for (int fontSize = maxFont; fontSize >= minFont; --fontSize) { |
| fontScale = fontSize / 20.0; |
| int base; |
| cv::Size textSize = |
| cv::getTextSize(text, fontFace, fontScale, thickness, &base); |
| if (textSize.width <= maxWidth) { |
| bestFontSize = fontSize; |
| outBaseline = base; |
| outFontFace = fontFace; |
| return fontScale; |
| } |
| } |
| outBaseline = 0; |
| outFontFace = fontFace; |
| return minFont / 20.0; |
| } |
|
|