diff --git a/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.cpp b/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.cpp index fbe0f98797..aa83363958 100644 --- a/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.cpp +++ b/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.cpp @@ -15,6 +15,10 @@ #include "ML/Models/ML_ONNXRuntimeHelpers.h" #include "ML_PaddleOCRPipeline.h" +#include +using std::cout; +using std::endl; + namespace PokemonAutomation{ namespace ML{ @@ -96,59 +100,49 @@ std::string PaddleOCRPipeline::recognize(const ImageViewRGB32& image){ // 1. Convert Image to OpenCV image (cv::mat) cv::Mat cv_image_rgb = imageviewrgb32_to_cv_mat_rgb(image); + if (cv_image_rgb.empty()) { + return ""; + } + - // 1b. Whitespace Trimming Logic - cv::Mat gray; - cv::cvtColor(cv_image_rgb, gray, cv::COLOR_BGR2GRAY); + // 2. Crop tightly around the text, with small safety margin + cv::Mat cropped_image = crop_to_text_region(cv_image_rgb); + if (cropped_image.empty()){ + return ""; + } - // Invert image for findNonZero: text becomes white (255), background black (0) - cv::Mat binary_inv; - cv::threshold(gray, binary_inv, 254, 255, cv::THRESH_BINARY_INV); + // add horizontal padding to tall/narrow characters + add_horizontal_padding(cropped_image); - // Find coordinates of all text pixels - std::vector nonZeroCoords; - cv::findNonZero(binary_inv, nonZeroCoords); - - cv::Mat cropped_image; - if (!nonZeroCoords.empty()){ - // Get the minimal bounding rectangle for the text - cv::Rect boundingBox = cv::boundingRect(nonZeroCoords); - - // Add a small buffer/padding around the box if needed (optional) - boundingBox.x = std::max(0, boundingBox.x - 2); - boundingBox.y = std::max(0, boundingBox.y - 2); - boundingBox.width = std::min(cv_image_rgb.cols - boundingBox.x, boundingBox.width + 4); - boundingBox.height = std::min(cv_image_rgb.rows - boundingBox.y, boundingBox.height + 4); - - // Crop the original color image - cropped_image = cv_image_rgb(boundingBox); - // static int i = 0; - // cv::imwrite("output" + std::to_string(i) + ".png", cropped_image); - // i++; - }else{ - return ""; // Return empty string if no text is detected in the region - } - - // 2a. Calculate dynamic width (maintain aspect ratio) + // 3. Calculate dynamic width (maintain aspect ratio) // the model shape is {1, 3, 48, dynamic_width}. Note that the height is fixed at 48 pixels // the input image must be scaled to match the height of 48, for the neural network int target_h = 48; - float aspect_ratio = (float)cropped_image.cols / (float)cropped_image.rows; - int target_w = static_cast(target_h * aspect_ratio); + float aspect_ratio = (float)cropped_image.cols / cropped_image.rows; + + int target_w = std::max( + 1, + (int)std::round(target_h * aspect_ratio) + ); - if (target_w <= 0) return ""; - - // 2b. Resize cv::Mat resized; - cv::resize(cropped_image, resized, cv::Size(target_w, target_h)); + cv::resize( + cropped_image, + resized, + cv::Size(target_w, target_h), + 0, + 0, + cv::INTER_LINEAR + ); - // 3. Normalize + + // 4. Normalize // convert UC3 8-bit [0,255] to 32FC3 float [0,1], then use ImageNet Normalization // output = (Input * Scale) = (old_pixel * 1/255). This transforms [0,255] to range [0, 1] // TODO: determine if normalizing to [-1,1] is preferred or to perform ImageNet normalization (mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225]) resized.convertTo(resized, CV_32FC3, 1.0 / 255.0); - // 3b. Apply Mean/Std (Standard for PaddleOCR). except for Chinese + // 4b. Apply Mean/Std (Standard for PaddleOCR). except for Chinese // Mean: [0.485, 0.456, 0.406], Std: [0.229, 0.224, 0.225] if (!(m_language == Language::ChineseSimplified || m_language == Language::ChineseTraditional || @@ -164,13 +158,13 @@ std::string PaddleOCRPipeline::recognize(const ImageViewRGB32& image){ } - // 3. Convert HWC to NCHW + // 5. Convert HWC to NCHW std::vector input_tensor_values = preprocess_NCHW(resized); - // 4. Define Dynamic Shape + // 6. Define Dynamic Shape std::vector input_shape = {1, 3, target_h, target_w}; - // 5. Create tensor with its own managed memory + // 7. Create tensor with its own managed memory Ort::AllocatorWithDefaultOptions allocator; auto input_tensor = Ort::Value::CreateTensor( allocator, input_shape.data(), input_shape.size() @@ -185,7 +179,7 @@ std::string PaddleOCRPipeline::recognize(const ImageViewRGB32& image){ const char* output_names[] = {m_output_name.c_str()}; try{ - // 7. Run the recognition session + // 8. Run the recognition session auto outputs = m_rec_session.Run( Ort::RunOptions{nullptr}, input_names, // char** @@ -201,6 +195,120 @@ std::string PaddleOCRPipeline::recognize(const ImageViewRGB32& image){ } +cv::Mat crop_to_text_region(const cv::Mat& image) { + // first convert to grayscale + cv::Mat gray; + cv::cvtColor(image, gray, cv::COLOR_RGB2GRAY); + + // get a binary image, for cropping purposes + cv::Mat binary; + cv::threshold(gray, binary, 0, 255, + cv::THRESH_BINARY_INV | cv::THRESH_OTSU); + + double ratio = cv::countNonZero(binary) / + static_cast(binary.total()); + + // we want text pixels to be white for the next step + // If most pixels are white, Otsu likely classified the background as foreground. + // Flip it so text becomes white again. + if (ratio > 0.5){ + cv::bitwise_not(binary, binary); + } + + // Find coordinates of all non-zero pixels (the text) + std::vector nonZeroCoords; + cv::findNonZero(binary, nonZeroCoords); + if (nonZeroCoords.empty()){ + return {}; + } + + // create bounding box for crop + cv::Rect bbox = cv::boundingRect(nonZeroCoords); + + // increase bounding box slightly to add small safety margin + int pad_x = std::max(4, bbox.width / 20); // ~5% + int pad_y = std::max(2, bbox.height / 20); // ~5% + + bbox.x = std::max(0, bbox.x - pad_x); + bbox.y = std::max(0, bbox.y - pad_y); + + bbox.width = std::min( + image.cols - bbox.x, + bbox.width + 2 * pad_x + ); + + bbox.height = std::min( + image.rows - bbox.y, + bbox.height + 2 * pad_y + ); + + // crop the original image based on the bounding box + // the crop should be within bounds. + cv::Mat cropped_image; + cropped_image = image(bbox).clone(); + + // static int i = 0; + // i++; + // cv::imwrite("aabinary" + std::to_string(i) + ".png", binary); + // cv::imwrite("aacropped_image" + std::to_string(i) + ".png", cropped_image); + + return cropped_image; +} + +void add_horizontal_padding(cv::Mat& image){ + if (image.empty()) { + return; + } + + int h = image.rows; + int w = image.cols; + constexpr float min_ratio = 0.5f; + + // add horizontal padding to tall/narrow characters + if (h > 0 && (float)w / h < min_ratio) { + int target_w = static_cast(std::ceil(min_ratio * h)); + + int left = (target_w - w) / 2; + int right = target_w - w - left; + + cv::Scalar bg = estimate_background_color(image); + cv::Mat padded_image; + cv::copyMakeBorder( + image, + padded_image, + 0, 0, // no vertical padding + left, right, + cv::BORDER_CONSTANT, + bg + ); + image = padded_image; + } + + // static int i = 0; + // i++; + // cv::imwrite("aapadded" + std::to_string(i) + ".png", image); + +} + +cv::Scalar estimate_background_color(const cv::Mat& image) { + if (image.empty() || image.type() != CV_8UC3) { + return cv::Scalar(255, 255, 255); + } + + // use the corner pixels to estimate the background color + const cv::Vec3b* top = image.ptr(0); + const cv::Vec3b* bottom = image.ptr(image.rows - 1); + + cv::Scalar tl(top[0]); + cv::Scalar tr(top[image.cols - 1]); + cv::Scalar bl(bottom[0]); + cv::Scalar br(bottom[image.cols - 1]); + + return (tl + tr + bl + br) * 0.25; +} + + + std::vector preprocess_NCHW(cv::Mat& img){ std::vector dst(img.rows * img.cols * 3); diff --git a/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.h b/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.h index 9fe2bbbdf2..3d7601ecf2 100644 --- a/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.h +++ b/SerialPrograms/Source/ML/Inference/ML_PaddleOCRPipeline.h @@ -48,6 +48,17 @@ class PaddleOCRPipeline { }; +// assumes the input image is RGB +cv::Mat crop_to_text_region(const cv::Mat& image); + +// if the image is narrow/tall, add horizontal padding +// modifies the input image +// assumes input image is RGB +void add_horizontal_padding(cv::Mat& image); + +// assumes input image is RGB +cv::Scalar estimate_background_color(const cv::Mat& image); + // convert HCW (height, width, channels) to NCHW (batch N, channels C, height H, width W) std::vector preprocess_NCHW(cv::Mat& img);