adjust material detector to use waterfill OCR instead of straight OCR.

PokemonAutomation · Jul 9, 2024 · 36dd97b · 36dd97b
1 parent af8eaff
commit 36dd97b
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 6 deletions.
diff --git a/SerialPrograms/Source/NintendoSwitch/DevPrograms/TestProgramSwitch.cpp b/SerialPrograms/Source/NintendoSwitch/DevPrograms/TestProgramSwitch.cpp
@@ -267,12 +267,52 @@ void TestProgram::program(MultiSwitchProgramEnvironment& env, CancellableScope&
 //    PokemonLA::ButtonDetector detector(logger, PokemonLA::ButtonType::ButtonA,);
 
 
+#if 1
+    ItemPrinterMaterialDetector detector(COLOR_RED, LANGUAGE);
+    detector.make_overlays(overlays);
+    // cout << (int)detector.find_happiny_dust_row_index(env.inference_dispatcher(), console, context) << endl;
+    // cout << (int)detector.detect_material_quantity(env.inference_dispatcher(), console, context, 2) << endl;
+
+    // test OCR for number 1 -> 999. for black text on light background.
+    // increasing quantity of materials to sell. 
+    for (int i = 1; i < 1000; i++){
+       context.wait_for_all_requests();
+        if (i != (int)detector.detect_material_quantity(env.inference_dispatcher(), console, context, 2)){
+            throw OperationFailedException(
+                ErrorReport::SEND_ERROR_REPORT,
+                console,
+                "OCR didn't match expected value."
+            );
+        }
+        pbf_press_dpad(context, DPAD_UP, 20, 30);
+    }
+
+    // test OCR for number 1 -> 999. for white text on dark background
+    // decreasing quantity of current materials by selling.
+    // for (int i = 999; i > 0; i--){
+    //     context.wait_for_all_requests();
+    //     if (i != (int)detector.detect_material_quantity(env.inference_dispatcher(), console, context, 2)){
+    //         throw OperationFailedException(
+    //             ErrorReport::SEND_ERROR_REPORT,
+    //             console,
+    //             "OCR didn't match expected value."
+    //         );
+    //     }
+    //     pbf_press_button(context, BUTTON_A, 30, 150);
+    //     pbf_press_button(context, BUTTON_A, 30, 150);
+    //     pbf_press_button(context, BUTTON_A, 30, 150);
+    //     pbf_press_button(context, BUTTON_A, 30, 150);
+    // }
 
+
+#endif
+
+#if 0
     ImageRGB32 image("screenshot-20240701-165012250266.png");
 
     BattleBallReader reader(console, Language::English);
     cout << reader.read_quantity(image) << endl;
-
+#endif
 
 #if 0
     ImageRGB32 image("screenshot-20240630-183016042676.png");

diff --git a/...Programs/Source/PokemonSV/Inference/ItemPrinter/PokemonSV_ItemPrinterMaterialDetector.cpp b/...Programs/Source/PokemonSV/Inference/ItemPrinter/PokemonSV_ItemPrinterMaterialDetector.cpp
@@ -11,6 +11,7 @@
 #include "CommonFramework/Exceptions/OperationFailedException.h"
 #include "CommonFramework/ImageTypes/ImageViewRGB32.h"
 #include "CommonFramework/ImageTypes/ImageRGB32.h"
+#include "CommonFramework/ImageTools/ImageStats.h"
 #include "CommonFramework/ImageTools/ImageFilter.h"
 #include "CommonFramework/OCR/OCR_NumberReader.h"
 #include "CommonFramework/Tools/ConsoleHandle.h"
@@ -84,13 +85,47 @@ int16_t ItemPrinterMaterialDetector::read_number(
     Logger& logger, AsyncDispatcher& dispatcher,
     const ImageViewRGB32& screen, const ImageFloatBox& box
 ) const{
-    int16_t number_result_white_text = read_number_black_or_white_text(logger, dispatcher, screen, box, true);
-    if (number_result_white_text == -1){
-        // try looking for black text
-        return read_number_black_or_white_text(logger, dispatcher, screen, box, false);
+
+    ImageViewRGB32 cropped = extract_box_reference(screen, box);
+    ImageRGB32 filtered = to_blackwhite_rgb32_range(
+        cropped,
+        0xff000000, 0xff808080,
+        true
+    );
+    // filtered.save("DebugDumps/test-one-filter-1.png");
+    bool is_dark_text_light_background = image_stats(filtered).average.sum() > 500;
+    // std::cout << "Average sum of filtered: "<< std::to_string(image_stats(filtered).average.sum()) << std::endl;
+
+    int16_t number;
+    if (is_dark_text_light_background){
+        number = (int16_t)OCR::read_number_waterfill(logger, cropped, 0xff000000, 0xff808080);
     }else{
-        return number_result_white_text;
+        number = (int16_t)OCR::read_number_waterfill(logger, cropped, 0xff808080, 0xffffffff);
+    }
+
+    if (number < 1 || number > 999){
+        number = -1;
     }
+    return (int16_t)number;
+
+    // return read_number_black_or_white_text(logger, dispatcher, screen, box, true);
+
+
+    // if (!is_dark_text_light_background){
+    //     filtered = to_blackwhite_rgb32_range(
+    //         extract_box_reference(screen, box),
+    //         0xff808080, 0xffffffff,
+    //         true
+    //     );
+    //     filtered.save("DebugDumps/test-one-filter-2.png");
+    // }
+
+    // ImageRGB32 scaled_image = filtered.scale_to(600, 200);
+    // scaled_image.save("DebugDumps/test-one-scaled-1.png");
+    // int value = OCR::read_number(logger, scaled_image);
+
+    // return (int16_t)value;
+
 
 }
 
@@ -143,6 +178,7 @@ int16_t ItemPrinterMaterialDetector::read_number_black_or_white_text(
     // }
     for (size_t c = 0; c < filtered.size(); c++){
         tasks[c] = dispatcher.dispatch([&, c]{
+            filtered[c].first.save("DebugDumps/test-" + std::to_string(c) + ".png");
             int num = OCR::read_number(logger, filtered[c].first);
             WriteSpinLock lg(lock);
             candidates[(int16_t)num]++;